Home
last modified time | relevance | path

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

/dports/math/py-jax/jax-0.2.9/jax/interpreters/
H A Dxla.py57 XlaShape = Any # xla_client.Shape variable
599 def flatten_shape(s: XlaShape) -> Sequence[Tuple[Sequence[int], XlaShape]]: argument
/dports/math/py-jax/jax-0.2.9/jax/experimental/
H A Dhost_callback.py379 XlaShape = Any # xla_client.Shape variable