Home
last modified time | relevance | path

Searched refs:jaxpr_uses_outfeed (Results 1 – 2 of 2) sorted by relevance

/dports/math/py-jax/jax-0.2.9/jax/interpreters/
H A Dxla.py202 def jaxpr_uses_outfeed(jaxpr: core.Jaxpr) -> bool: function
209 if jaxpr_uses_outfeed(param):
212 if jaxpr_uses_outfeed(param.jaxpr):
/dports/math/py-jax/jax-0.2.9/jax/experimental/
H A Dhost_callback.py1097 if not has_input_token and not xla.jaxpr_uses_outfeed(jaxpr):
1155 if xla.jaxpr_uses_outfeed(cond_jaxpr.jaxpr):