Filtering ops for a backend. #8995
-
|
Hi JAX team, This question is from the perspective of adding a new backend. In the TF/XLA interface, it is possible to specify a callback to filter supported operations when registering a device. For example: Is there something similar in the JAX/XLA interface? In the backend registration I did not see something similar: https://github.com/google/jax/blob/main/jax/_src/lib/xla_bridge.py#L200 Thanks! |
Beta Was this translation helpful? Give feedback.
Replies: 1 comment
-
|
There's no exact equivalent. We assume all backends can implement all operations. If that's a difficult assumption, we'd need to discuss the use case in a bit more detail. What you can do is define platform-specific XLA lowerings for a rule, e.g., see how we call LAPACK or Cusolver. |
Beta Was this translation helpful? Give feedback.
There's no exact equivalent. We assume all backends can implement all operations. If that's a difficult assumption, we'd need to discuss the use case in a bit more detail.
What you can do is define platform-specific XLA lowerings for a rule, e.g., see how we call LAPACK or Cusolver.