That’s an interesting idea and could definitely be done. You already identified that we could probably just copy a lot from the JAX implementation. This is probably better discussed as an aesara issue.
That’s an interesting idea and could definitely be done. You already identified that we could probably just copy a lot from the JAX implementation. This is probably better discussed as an aesara issue.