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

Fix accidental signature change to get_serialized_metadata() from nanobind PR.

Browse Source 
pybind11 accepts either Python strings or bytes as a std::string argument, whereas nanobind accepts only strings. Change the argument to nb::bytes instead.

PiperOrigin-RevId: 560086072
This commit is contained in:
Peter Hawkins 2023-08-25 07:30:44 -07:00 committed by jax authors
parent 87cec1a6e8
commit ac8ea86103
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
jaxlib/gpu/triton.cc
View File

@ -134,10 +134,11 @@ NB_MODULE(_triton, m) {
m.def("get_serialized_metadata",
ValueOrThrowWrapper(
[](std::string_view opaque) -> absl::StatusOr<nb::bytes> {
[](nb::bytes opaque) -> absl::StatusOr<nb::bytes> {
JAX_ASSIGN_OR_RETURN(
std::string metadata,
GetTritonKernelCallSerializedMetadata(opaque));
GetTritonKernelCallSerializedMetadata(
absl::string_view(opaque.c_str(), opaque.size())));
return nb::bytes(metadata.c_str(), metadata.size());
}));
}