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

rm CmdBuffer traits

Browse Source 
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
This commit is contained in:
Phuong Nguyen 2024-10-16 10:27:09 -07:00
parent f3775aa233
commit 82113cd047
2 changed files with 2 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
docs/cuda_custom_call/cuda_custom_call_test.py
View File

@ -27,7 +27,6 @@ import numpy as np
import jax
import jax.numpy as jnp
from jax.extend import ffi
from jax.lib import xla_client
# start test boilerplate
from absl.testing import absltest

6
docs/cuda_custom_call/foo.cu.cc
View File

@ -77,8 +77,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(
.Arg<ffi::Buffer<ffi::DataType::F32>>() // b
.Ret<ffi::Buffer<ffi::DataType::F32>>() // c
.Ret<ffi::Buffer<ffi::DataType::F32>>() // b_plus_1
.Attr<size_t>("n"),
{xla::ffi::Traits::kCmdBufferCompatible}); // cudaGraph enabled
.Attr<size_t>("n"));
//----------------------------------------------------------------------------//
// Backward pass //
@ -136,5 +135,4 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(
.Arg<ffi::Buffer<ffi::DataType::F32>>() // b_plus_1
.Ret<ffi::Buffer<ffi::DataType::F32>>() // a_grad
.Ret<ffi::Buffer<ffi::DataType::F32>>() // b_grad
.Attr<size_t>("n"),
{xla::ffi::Traits::kCmdBufferCompatible}); // cudaGraph enabled
.Attr<size_t>("n"));