mirror of
https://github.com/ROCm/jax.git
synced 2025-04-18 21:06:06 +00:00
83ffcc9c7d
Add print First version with custom_partitioning. The communication during the gradient aren't optimal. Fix the gradient sharding small update Fix the strange replicated computation. Make it work with the new JAX version. Add the structure for custom_p domentation. Small clean up First version of the doc Add comment and typing annotation tab->space Simplify code and add docstring Use the simpler JAX API since 0.4.16 (August 2023). Custom partitioning using custom_partitioning updated docs; dump custom_partitioning HLO doc update more documentation updates; include links to code instead of inlined code fix typos fix more typos fix type annotations in source and update docs minor fixes import fix lint fix added apache license header
42 lines
1.5 KiB
C++
42 lines
1.5 KiB
C++
/* Copyright 2024 The JAX Authors.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
==============================================================================*/
|
|
|
|
// This header extends kernel_helpers.h with the pybind11 specific interface to
|
|
// serializing descriptors. It also adds a pybind11 function for wrapping our
|
|
// custom calls in a Python capsule. This is separate from kernel_helpers so
|
|
// that the CUDA code itself doesn't include pybind11. I don't think that this
|
|
// is strictly necessary, but they do it in jaxlib, so let's do it here too.
|
|
|
|
#ifndef _GPU_OPS_PYBIND11_KERNEL_HELPERS_H_
|
|
#define _GPU_OPS_PYBIND11_KERNEL_HELPERS_H_
|
|
|
|
#include <pybind11/pybind11.h>
|
|
|
|
#include "kernel_helpers.h"
|
|
|
|
namespace gpu_ops {
|
|
|
|
template <typename T> pybind11::bytes PackDescriptor(const T &descriptor) {
|
|
return pybind11::bytes(PackDescriptorAsString(descriptor));
|
|
}
|
|
|
|
template <typename T> pybind11::capsule EncapsulateFunction(T *fn) {
|
|
return pybind11::capsule(bit_cast<void *>(fn), "xla._CUSTOM_CALL_TARGET");
|
|
}
|
|
|
|
} // namespace gpu_ops
|
|
|
|
#endif
|