Home
last modified time | relevance | path

Searched defs:ShapedArray (Results 1 – 5 of 5) sorted by relevance

/dports/math/py-jax/jax-0.2.9/jax/
H A Dabstract_arrays.py29 ShapedArray = core.ShapedArray variable
H A Dcore.py1030 class ShapedArray(UnshapedArray): class
/dports/math/py-jax/jax-0.2.9/jax/interpreters/
H A Dxla.py114 def array_result_handler(device: Optional[Device], aval: core.ShapedArray):
1029 aval: core.ShapedArray,
1071 def __init__(self, aval: core.ShapedArray, device: Optional[Device],
1349 def _lazy_force_computation(aval: core.ShapedArray,
H A Dpxla.py467 def array_result_handler(sharding_spec, indices, aval: ShapedArray):
505 aval: ShapedArray,
/dports/math/py-jax/jax-0.2.9/jax/_src/lax/
H A Dlax.py2666 lhs: ShapedArray, rhs: ShapedArray, *, window_strides, padding,