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

Add padding rule for copy_p

Browse Source
This commit is contained in:
Jake VanderPlas 2022-08-19 14:10:12 -07:00
parent 314cf8a439
commit 1d920b51a9
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

@ -4318,6 +4318,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)