Bend Waveguide

import jax.numpy as jnp
from rich import print as rprint
import sax