File tree Expand file tree Collapse file tree 2 files changed +6
-15
lines changed Expand file tree Collapse file tree 2 files changed +6
-15
lines changed Original file line number Diff line number Diff line change @@ -9,6 +9,11 @@ namespace py = pybind11;
99namespace infinicore ::ops {
1010
1111inline void bind_rope (py::module &m) {
12+
13+ py::enum_<infinicore::nn::RoPE::Algo>(m, " RoPEAlgo" )
14+ .value (" GPT_J" , infinicore::nn::RoPE::Algo::GPT_J)
15+ .value (" GPT_NEOX" , infinicore::nn::RoPE::Algo::GPT_NEOX);
16+
1217 m.def (" rope" ,
1318 &op::rope,
1419 py::arg (" x" ),
Original file line number Diff line number Diff line change @@ -22,21 +22,7 @@ void TensorImpl::copy_from(Tensor src) {
2222 throw std::runtime_error (" Cannot copy from tensor with different shape" );
2323 }
2424 if (this ->device () == src->device ()) {
25-
26- // If both tensors are contiguous, use direct memcpy (much faster and avoids rearrange issues)
27- if (this ->is_contiguous () && src->is_contiguous ()) {
28- // Use nbytes() to get the actual tensor size
29- size_t copy_size = std::min (this ->nbytes (), src->nbytes ());
30-
31- // For CPU-to-CPU copies, use regular memcpy. For device-to-device, use D2D memcpy
32- if (this ->device ().getType () == Device::Type::CPU) {
33- context::memcpyH2H (this ->data (), src->data (), copy_size);
34- } else {
35- context::memcpyD2D (this ->data (), src->data (), copy_size);
36- }
37- } else {
38- op::rearrange_ (Tensor (const_cast <TensorImpl *>(this )->shared_from_this ()), src);
39- }
25+ op::rearrange_ (Tensor (const_cast <TensorImpl *>(this )->shared_from_this ()), src);
4026 } else {
4127 if (!src->is_contiguous ()) {
4228 src = src->contiguous ();
You can’t perform that action at this time.
0 commit comments