| /dports/math/py-jax/jax-0.2.9/jax/experimental/ |
| H A D | maps.py | 187 in_axes, argument
267 def xmap_impl(fun: lu.WrappedFun, *args, name, in_axes, out_axes_thunk, axis_sizes, argument
276 in_axes, out_axes_thunk, axis_sizes, argument
324 def vectorize(self, f: lu.WrappedFun, in_axes, out_axes): argument
332 def to_mesh_axes(self, in_axes, out_axes): argument
453 in_axes, out_axes, axis_sizes, argument
517 def _xla_tile(c, axis_env, x, in_axes, axis_sizes): argument
560 in_axes, out_axes, axis_sizes, argument
|
| /dports/math/py-flax/flax-0.3.3/flax/linen/ |
| H A D | transforms.py | 215 in_axes=0, out_axes=0, argument
370 in_axes=0, out_axes=0, argument
|
| /dports/math/py-jax/jax-0.2.9/jax/interpreters/ |
| H A D | pxla.py | 640 global_axis_size, devices, name, in_axes, out_axes_thunk, argument
1221 call_jaxpr, *, backend=None, in_axes, out_axes, argument
1603 def soft_pmap_impl(fun: lu.WrappedFun, *args, axis_name, axis_size, in_axes, out_axes_thunk): argument
1610 def _soft_pmap_callable(fun, axis_name, axis_size, in_axes, out_axes_thunk, *avals): argument
1665 def _soft_pmap_jaxpr(jaxpr, consts, in_axes, axis_name, axis_size, chunk_size): argument
|
| /dports/math/py-flax/flax-0.3.3/flax/core/ |
| H A D | lift.py | 313 in_axes=0, out_axes=0, argument
421 in_axes=0, out_axes=0, argument
|
| /dports/math/py-jax/jax-0.2.9/jax/ |
| H A D | api.py | 1073 def vmap(fun: F, in_axes=0, out_axes=0, axis_name=None) -> F: argument
1282 in_axes=0, argument
1580 def soft_pmap(fun: Callable, axis_name: Optional[AxisName] = None, in_axes=0 argument
|
| /dports/misc/mxnet/incubator-mxnet-1.9.0/3rdparty/tvm/src/relay/op/tensor/ |
| H A D | reduce.cc | 58 std::vector<int64_t> in_axes; in GetReduceAxes() local
|
| /dports/misc/tvm/incubator-tvm-0.6.1/src/relay/op/tensor/ |
| H A D | reduce.cc | 55 std::vector<int64_t> in_axes; in GetReduceAxes() local
|
| /dports/misc/py-tvm/incubator-tvm-0.6.1/src/relay/op/tensor/ |
| H A D | reduce.cc | 55 std::vector<int64_t> in_axes; in GetReduceAxes() local
|
| /dports/math/py-matplotlib2/matplotlib-2.2.4/lib/matplotlib/axes/ |
| H A D | _base.py | 2197 def in_axes(self, mouseevent): member in _AxesBase
|
| /dports/math/py-matplotlib/matplotlib-3.4.3/lib/matplotlib/axes/ |
| H A D | _base.py | 2362 def in_axes(self, mouseevent): member in _AxesBase
|