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

Merge pull request #12021 from jakevdp:copy-padding

Browse Source 
PiperOrigin-RevId: 469292408
This commit is contained in:
jax authors 2022-08-22 15:02:39 -07:00
commit 3aa42775e7
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
jax/_src/lax/lax.py
View File

@ -4321,6 +4321,7 @@ copy_p.def_impl(partial(xla.apply_primitive, copy_p))
copy_p.def_abstract_eval(lambda x: x)
mlir.register_lowering(copy_p, lambda ctx, x: [x])
ad.deflinear(copy_p, lambda t: [copy_p.bind(t)])
pe.def_trivial_padding(copy_p)
batching.defvectorized(copy_p)