adding __float/bool/int___

Author: fschlimb

ddptensor/array_api.py

@@ -144,3 +144,13 @@
     "std",   # (x, /, *, axis=None, correction=0.0, keepdims=False)
     "var",   # (x, /, *, axis=None, correction=0.0, keepdims=False)
 ]
+
+unary_methods = [
+    "__array_namespace__",  # (self, /, *, api_version=None)
+    "__bool__",  # (self, /)
+    "__dlpack__",  # (self, /, *, stream=None)
+    "__dlpack_device__",  # (self, /)
+    "__float__",  # (self, /)
+    "__int__",  # (self, /)
+    "__len__",  # (self, /)
+]

ddptensor/ddptensor.py

@@ -2,16 +2,6 @@
 from ._ddptensor import float64, int64, fini
 from . import array_api as api
 
-unary_methods = [
-    # "__array_namespace__",  # (self, /, *, api_version=None)
-    "__bool__",  # (self, /)
-    # "__dlpack__",  # (self, /, *, stream=None)
-    # "__dlpack_device__",  # (self, /)
-    "__float__",  # (self, /)
-    "__int__",  # (self, /)
-    "__len__",  # (self, /)
-]
-
 t_attributes = ["dtype", "shape", ]  #"device", "ndim", "size", "T"]
 
 #def try_except(func, *args, **kwargs):
@@ -27,7 +17,6 @@ def __init__(self, t):
     def __repr__(self):
         return self._t.__repr__()
 
-
     for method in api.ew_binary_methods:
         METHOD = method.upper()
         exec(
@@ -45,11 +34,11 @@ def __repr__(self):
         exec(
             f"{method} = lambda self: dtensor(_cdt.EWUnyOp.op(_cdt.{METHOD}, self._t))"
         )
-        
-    # for method in unary_methods:
-    #     exec(
-    #         f"{method} = lambda self: self._t.{method}()"
-    #     )
+
+    for method in api.unary_methods:
+        exec(
+            f"{method} = lambda self: _cdt.UnyOp.{method}(self._t)"
+        )
 
     # for att in t_attributes:
     #     exec(

src/ddptensor.cpp

@@ -92,6 +92,24 @@ struct EWUnyOp
     }
 };
 
+struct UnyOp
+{
+    static bool __bool__(x::DPTensorBaseX::ptr_type a)
+    {
+        return TypeDispatch<x::UnyOp>(a->dtype(), a, true);
+    }
+
+    static double __float__(x::DPTensorBaseX::ptr_type a)
+    {
+        return TypeDispatch<x::UnyOp>(a->dtype(), a, double(1));
+    }
+
+    static int64_t __int__(x::DPTensorBaseX::ptr_type a)
+    {
+        return TypeDispatch<x::UnyOp>(a->dtype(), a, int64_t(1));
+    }
+};
+
 struct ReduceOp
 {
     static auto op(ReduceOpId op, x::DPTensorBaseX::ptr_type a, const dim_vec_type & dim)
@@ -148,6 +166,11 @@ PYBIND11_MODULE(_ddptensor, m) {
     py::class_<EWUnyOp>(m, "EWUnyOp")
         .def("op", &EWUnyOp::op);
 
+    py::class_<UnyOp>(m, "UnyOp")
+        .def("__bool__", &UnyOp::__bool__)
+        .def("__float__", &UnyOp::__float__)
+        .def("__int__", &UnyOp::__int__);
+
     py::class_<IEWBinOp>(m, "IEWBinOp")
         .def("op", &IEWBinOp::op);
 

src/include/ddptensor/PVSlice.hpp

@@ -155,6 +155,11 @@ class PVSlice
         return _slice;
     }
 
+    uint64_t size() const
+    {
+        return slice().size();
+    }
+
 #if 0
     NDSlice normalized_slice() const
     {

src/include/ddptensor/x.hpp

@@ -47,26 +47,37 @@ namespace x
     template<typename T>
     class DPTensorX : public DPTensorBaseX
     {
+        rank_type _owner;
         PVSlice _slice;
         xt::xstrided_slice_vector _lslice;
         std::shared_ptr<xt::xarray<T>> _xarray;
+        T _replica = 0;
 
     public:
         template<typename I>
-        DPTensorX(PVSlice && slc, I && ax)
-            : _slice(std::move(slc)),
+        DPTensorX(PVSlice && slc, I && ax, rank_type owner=NOOWNER)
+            : _owner(owner),
+              _slice(std::move(slc)),
               _lslice(to_xt(_slice.local_slice_of_rank())),
               _xarray(std::make_shared<xt::xarray<T>>(std::forward<I>(ax)))
         {
         }
 
         template<typename O>
-        DPTensorX(const DPTensorX<O> & org, const NDSlice & slc)
-            : _slice(org._slice, slc),
+        DPTensorX(const DPTensorX<O> & org, const NDSlice & slc, rank_type owner=NOOWNER)
+            : _owner(owner),
+              _slice(org._slice, slc),
               _lslice(to_xt(_slice.local_slice_of_rank())),
               _xarray(org._xarray)
         {
-            std::cerr << "slice: " << _slice.slice() << " lslice: " << _slice.local_slice_of_rank() << std::endl;
+            if(owner == NOOWNER && slice().size() <= 1) {
+                set_owner(org.slice().owner(slc));
+            } else if(owner == REPLICATED) {
+                _replica = *(xt::strided_view(xarray(), to_xt(slice().slice())).begin());
+            }
+            std::cerr << "slice: " << _slice.slice() << " sz " << _slice.size()
+                      << " lslice: " << _slice.local_slice_of_rank() << " owner: " << _owner
+                      << " val: " << _replica << std::endl;
         }
 
         virtual std::string __repr__() const
@@ -82,6 +93,11 @@ namespace x
             return DTYPE<T>::value;
         }
 
+        virtual shape_type shape() const
+        {
+            return _slice.shape();
+        }
+
         xt::xarray<T> & xarray()
         {
             return *_xarray.get();
@@ -102,9 +118,41 @@ namespace x
             return _lslice;
         }
 
-        virtual shape_type shape() const
+        bool has_owner() const
         {
-            return _slice.shape();
+            return _owner < _OWNER_END;
+        }
+        
+        void set_owner(rank_type o)
+        {
+            _owner = o;
+        }
+        
+        rank_type owner() const
+        {
+            return _owner;
+        }
+        
+        bool is_replicated() const
+        {
+            return _owner == REPLICATED;
+        }
+
+        T replicate()
+        {
+            std::cerr << "is_replicated()=" << is_replicated() << " owner=" << owner() << " shape=" << to_string(shape()) << std::endl;
+            if(is_replicated()) return _replica;
+            if(has_owner() && _slice.size() == 1) {
+                if(theTransceiver->rank() == owner()) {
+                    _replica = *(xt::strided_view(xarray(), lslice()).begin());
+                    std::cerr << "replica: " << _replica << std::endl;
+                }
+                theTransceiver->bcast(&_replica, sizeof(T), owner());
+                set_owner(REPLICATED);
+            } else {
+                throw(std::runtime_error("Replication implemented for single element and single owner only."));
+            }
+            return _replica;
         }
     };
 
@@ -460,6 +508,34 @@ namespace x
 
     };
 
+    template<typename T>
+    class UnyOp
+    {
+    public:
+        using ptr_type = DPTensorBaseX::ptr_type;
+
+        template<typename N>
+        static N __type__(const ptr_type & a_ptr)
+        {
+            auto const _a = dynamic_cast<DPTensorX<T>*>(a_ptr.get());
+            if(!_a )
+                throw std::runtime_error("Invalid array object: could not dynamically cast");
+            T v = _a->replicate();
+            return static_cast<N>(v);
+        }
+        static bool op(const ptr_type & a_ptr, bool)
+        {
+            return __type__<bool>(a_ptr);
+        }
+        static double op(const ptr_type & a_ptr, double)
+        {
+            return __type__<double>(a_ptr);
+        }
+        static int64_t op(const ptr_type & a_ptr, int64_t)
+        {
+            return __type__<int64_t>(a_ptr);
+        }
+    };
 
     template<typename T>
     class ReduceOp
@@ -474,11 +550,13 @@ namespace x
         {
             xt::xarray<typename X::value_type> a = x;
             auto new_shape = reduce_shape(slice.shape(), dims);
+            rank_type owner = NOOWNER;
             if(slice.need_reduce(dims)) {
                 auto len = VPROD(new_shape);
                 theTransceiver->reduce_all(a.data(), DTYPE<typename X::value_type>::value, len, rop);
+                owner = REPLICATED;
             }
-            return std::make_shared<DPTensorX<typename X::value_type>>(new_shape, a);
+            return std::make_shared<DPTensorX<typename X::value_type>>(new_shape, a, owner);
         }
 
         static ptr_type op(ReduceOpId rop, const ptr_type & a_ptr, const dim_vec_type & dims)
@@ -529,5 +607,5 @@ namespace x
             return std::make_shared<DPTensorX<T>>(*_a, slice);
         }
     };
-   
+
 } // namespace x

test/test_x.py

@@ -6,5 +6,6 @@
 print(a == b)
 print(dt.sqrt(a))
 print(dt.sum(a, [1]))
-print(a[0:1,0:1])
+print(a[0:1,0:1], float(a[0:1,0:1]), bool(a[0:1,0:1]), int(a[0:1,0:1]))
+print(a[0:2,0:2])
 dt.fini()