mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 11:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #14788 from patrick-kidger:patch-1

Browse Source 
PiperOrigin-RevId: 514437372
This commit is contained in:
jax authors 2023-03-06 09:40:19 -08:00
commit 0d24d79453
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/_src/interpreters/batching.py
View File

@ -815,7 +815,7 @@ def _matchaxis_symbolic_zeros(axis_name, sz, name, src, dst, x, sum_match=False)
### utilities for defining primitives' batching rules
BatchingRule = Callable[..., Tuple[Any, Union[int, Tuple[int, ...]]]]
BatchingRule = Callable[..., Tuple[Any, Union[None, int, Tuple[Union[None, int], ...]]]]
primitive_batchers : Dict[core.Primitive, BatchingRule] = {}
axis_primitive_batchers: Dict[core.Primitive, Callable] = {}
spmd_axis_primitive_batchers: Dict[core.Primitive, Callable] = {}