From 53d4ea1f5256bbca2430f241a0aec9afe93b5991 Mon Sep 17 00:00:00 2001 From: Matthew Parno Date: Tue, 27 Feb 2024 09:08:04 -0500 Subject: [PATCH 1/6] Fixed pickling of pytorch wrapper. --- bindings/python/package/torch.py | 120 ++------------------- bindings/python/package/torch_helpers.py | 115 ++++++++++++++++++++ bindings/python/tests/test_TorchWrapper.py | 35 ++++-- 3 files changed, 147 insertions(+), 123 deletions(-) create mode 100644 bindings/python/package/torch_helpers.py diff --git a/bindings/python/package/torch.py b/bindings/python/package/torch.py index 166a5ecc..d2a1334e 100644 --- a/bindings/python/package/torch.py +++ b/bindings/python/package/torch.py @@ -1,115 +1,7 @@ import torch +import dill -def ExtractTorchTensorData(tensor): - """ Extracts the pointer, shape, and stride from a pytorch tensor and returns a tuple - that can be passed to MParT functions that have been overloaded to accept - (double*, std::tuple, std::tuple) instead of a Kokkos::View. - - Arguments: - ------------ - tensor: pytorch.Tensor - The pytorch tensor we want to eventually wrap with a Kokkos view. - - Returns: - ------------ - Tuple[int, Tuple[int,int], Tuple[int,int]] - A python tuple that contains all information needed to construct a Kokkos::View. - After casting to c++ types using pybind, this output can be passed to the - mpart::ConstructViewFromPointer function. - """ - - # Make sure the tensor has double data type - if tensor.dtype != torch.float64: - raise ValueError(f'Currently only tensors with float64 datatype can be converted. Current dtype is {tensor.dtype}') - - if len(tensor.shape)==1: - return tensor.data_ptr(), tensor.shape[0], tensor.stride()[0] - elif len(tensor.shape)==2: - return tensor.data_ptr(), tuple(tensor.shape), tuple(tensor.stride()) - else: - raise ValueError(f'Currently only 1d and 2d tensors can be converted.') - - -class MpartTorchAutograd(torch.autograd.Function): - - @staticmethod - def forward(ctx, input, coeffs, f, return_logdet): - ctx.save_for_backward(input, coeffs) - ctx.f = f - - coeffs_dbl = None - if coeffs is not None: - coeffs_dbl = coeffs.double() - f.WrapCoeffs(ExtractTorchTensorData(coeffs_dbl)) - input_dbl = input.double() - - output = torch.zeros(f.outputDim, input.shape[1], dtype=torch.double) - f.EvaluateImpl(ExtractTorchTensorData(input_dbl), ExtractTorchTensorData(output)) - - if return_logdet: - logdet = torch.zeros(input.shape[1], dtype=torch.double) - f.LogDeterminantImpl(ExtractTorchTensorData(input_dbl), ExtractTorchTensorData(logdet)) - return output.type(input.dtype), logdet.type(input.dtype) - else: - return output.type(input.dtype) - - @staticmethod - def backward(ctx, output_sens, logdet_sens=None): - input, coeffs = ctx.saved_tensors - f = ctx.f - - coeffs_dbl = None - if coeffs is not None: - coeffs_dbl = coeffs.double() - f.WrapCoeffs(ExtractTorchTensorData(coeffs_dbl)) - input_dbl = input.double() - output_sens_dbl = output_sens.double() - - logdet_sens_dbl = None - if logdet_sens is not None: - logdet_sens_dbl = logdet_sens.double() - - # Get the gradient wrt input - grad = None - if input.requires_grad: - grad = torch.zeros(f.inputDim, input.shape[1], dtype=torch.double) - - f.GradientImpl(ExtractTorchTensorData(input_dbl), - ExtractTorchTensorData(output_sens_dbl), - ExtractTorchTensorData(grad)) - - if logdet_sens is not None: - grad2 = torch.zeros(f.inputDim, input.shape[1], dtype=torch.double) - - f.LogDeterminantInputGradImpl(ExtractTorchTensorData(input_dbl), - ExtractTorchTensorData(grad2)) - grad += grad2*logdet_sens_dbl[None,:] - - coeff_grad = None - if coeffs is not None: - if coeffs.requires_grad: - coeff_grad = torch.zeros(f.numCoeffs, input.shape[1], dtype=torch.double) - f.CoeffGradImpl(ExtractTorchTensorData(input_dbl), - ExtractTorchTensorData(output_sens_dbl), - ExtractTorchTensorData(coeff_grad)) - - coeff_grad = coeff_grad.sum(axis=1) # pytorch expects total gradient not per-sample gradient - - if logdet_sens is not None: - grad2 = torch.zeros(f.numCoeffs, input.shape[1], dtype=torch.double) - - f.LogDeterminantCoeffGradImpl(ExtractTorchTensorData(input_dbl), - ExtractTorchTensorData(grad2)) - - coeff_grad += torch.sum(grad2*logdet_sens[None,:],axis=1) - - if coeff_grad is not None: - coeff_grad = coeff_grad.type(input.dtype) - - if grad is not None: - grad = grad.type(input.dtype) - - return grad, coeff_grad, None, None +from .torch_helpers import ExtractTorchTensorData, MpartTorchAutograd @@ -118,7 +10,7 @@ class TorchParameterizedFunctionBase(torch.nn.Module): can be used with pytorch. """ - def __init__(self, f, store_coeffs=True, dtype=torch.double): + def __init__(self, f=None, store_coeffs=True, dtype=torch.double): super().__init__() self.f = f @@ -129,7 +21,7 @@ def __init__(self, f, store_coeffs=True, dtype=torch.double): self.coeffs = torch.nn.Parameter(coeff_tensor) else: self.coeffs = None - + def forward(self, x, coeffs=None): if coeffs is None: @@ -148,7 +40,7 @@ class TorchConditionalMapBase(torch.nn.Module): This can be done either in the constructor or afterwards. """ - def __init__(self, f, store_coeffs=True, return_logdet=False, dtype=torch.double): + def __init__(self, f=None, store_coeffs=True, return_logdet=False, dtype=torch.double): super().__init__() self.return_logdet = return_logdet @@ -159,7 +51,7 @@ def __init__(self, f, store_coeffs=True, return_logdet=False, dtype=torch.double self.coeffs = torch.nn.Parameter(coeff_tensor) else: self.coeffs = None - + def forward(self, x, coeffs=None): if coeffs is None: diff --git a/bindings/python/package/torch_helpers.py b/bindings/python/package/torch_helpers.py new file mode 100644 index 00000000..8fe5cfbf --- /dev/null +++ b/bindings/python/package/torch_helpers.py @@ -0,0 +1,115 @@ +import torch + +def ExtractTorchTensorData(tensor): + """ Extracts the pointer, shape, and stride from a pytorch tensor and returns a tuple + that can be passed to MParT functions that have been overloaded to accept + (double*, std::tuple, std::tuple) instead of a Kokkos::View. + + Arguments: + ------------ + tensor: pytorch.Tensor + The pytorch tensor we want to eventually wrap with a Kokkos view. + + Returns: + ------------ + Tuple[int, Tuple[int,int], Tuple[int,int]] + A python tuple that contains all information needed to construct a Kokkos::View. + After casting to c++ types using pybind, this output can be passed to the + mpart::ConstructViewFromPointer function. + """ + + # Make sure the tensor has double data type + if tensor.dtype != torch.float64: + raise ValueError(f'Currently only tensors with float64 datatype can be converted. Current dtype is {tensor.dtype}') + + if len(tensor.shape)==1: + return tensor.data_ptr(), tensor.shape[0], tensor.stride()[0] + elif len(tensor.shape)==2: + return tensor.data_ptr(), tuple(tensor.shape), tuple(tensor.stride()) + else: + raise ValueError(f'Currently only 1d and 2d tensors can be converted.') + + +class MpartTorchAutograd(torch.autograd.Function): + + def __reduce__(self): + return (self.__class__, (None,)) + + @staticmethod + def forward(ctx, input, coeffs, f, return_logdet): + ctx.save_for_backward(input, coeffs) + ctx.f = f + + coeffs_dbl = None + if coeffs is not None: + coeffs_dbl = coeffs.double() + f.WrapCoeffs(ExtractTorchTensorData(coeffs_dbl)) + input_dbl = input.double() + + output = torch.zeros(f.outputDim, input.shape[1], dtype=torch.double) + f.EvaluateImpl(ExtractTorchTensorData(input_dbl), ExtractTorchTensorData(output)) + + if return_logdet: + logdet = torch.zeros(input.shape[1], dtype=torch.double) + f.LogDeterminantImpl(ExtractTorchTensorData(input_dbl), ExtractTorchTensorData(logdet)) + return output.type(input.dtype), logdet.type(input.dtype) + else: + return output.type(input.dtype) + + @staticmethod + def backward(ctx, output_sens, logdet_sens=None): + input, coeffs = ctx.saved_tensors + f = ctx.f + + coeffs_dbl = None + if coeffs is not None: + coeffs_dbl = coeffs.double() + f.WrapCoeffs(ExtractTorchTensorData(coeffs_dbl)) + input_dbl = input.double() + output_sens_dbl = output_sens.double() + + logdet_sens_dbl = None + if logdet_sens is not None: + logdet_sens_dbl = logdet_sens.double() + + # Get the gradient wrt input + grad = None + if input.requires_grad: + grad = torch.zeros(f.inputDim, input.shape[1], dtype=torch.double) + + f.GradientImpl(ExtractTorchTensorData(input_dbl), + ExtractTorchTensorData(output_sens_dbl), + ExtractTorchTensorData(grad)) + + if logdet_sens is not None: + grad2 = torch.zeros(f.inputDim, input.shape[1], dtype=torch.double) + + f.LogDeterminantInputGradImpl(ExtractTorchTensorData(input_dbl), + ExtractTorchTensorData(grad2)) + grad += grad2*logdet_sens_dbl[None,:] + + coeff_grad = None + if coeffs is not None: + if coeffs.requires_grad: + coeff_grad = torch.zeros(f.numCoeffs, input.shape[1], dtype=torch.double) + f.CoeffGradImpl(ExtractTorchTensorData(input_dbl), + ExtractTorchTensorData(output_sens_dbl), + ExtractTorchTensorData(coeff_grad)) + + coeff_grad = coeff_grad.sum(axis=1) # pytorch expects total gradient not per-sample gradient + + if logdet_sens is not None: + grad2 = torch.zeros(f.numCoeffs, input.shape[1], dtype=torch.double) + + f.LogDeterminantCoeffGradImpl(ExtractTorchTensorData(input_dbl), + ExtractTorchTensorData(grad2)) + + coeff_grad += torch.sum(grad2*logdet_sens[None,:],axis=1) + + if coeff_grad is not None: + coeff_grad = coeff_grad.type(input.dtype) + + if grad is not None: + grad = grad.type(input.dtype) + + return grad, coeff_grad, None, None \ No newline at end of file diff --git a/bindings/python/tests/test_TorchWrapper.py b/bindings/python/tests/test_TorchWrapper.py index 70dbc4d5..5c28a1a8 100644 --- a/bindings/python/tests/test_TorchWrapper.py +++ b/bindings/python/tests/test_TorchWrapper.py @@ -7,6 +7,7 @@ import mpart as mt import numpy as np +import dill if haveTorch: @@ -81,6 +82,21 @@ def test_AutogradCoeffs(): loss.backward() assert tmap2.coeffs.grad is not None + def test_AutogradCoeffAsInput(): + + opts = mt.MapOptions() + tmap = mt.CreateTriangular(dim,dim,3,opts) # Simple third order map + + tmap2 = tmap.torch(store_coeffs=False) + + x = torch.randn(numSamps, dim, dtype=torch.double) + coeffs = torch.randn(tmap.numCoeffs, dtype=torch.double) + + y = tmap2(x,coeffs) + assert y.shape[0] == numSamps + assert y.shape[1] == dim + assert not y.isnan().any() + def test_TorchMethod(): opts = mt.MapOptions() tmap = mt.CreateTriangular(dim,dim,3,opts) # Simple third order map @@ -96,20 +112,20 @@ def test_TorchMethod(): assert np.all(y.detach().numpy() == tmap.Evaluate(x.T.detach().numpy()).T) assert np.all(logdet.detach().numpy() == tmap.LogDeterminant(x.T.detach().numpy())) - def test_AutogradCoeffAsInput(): - + def test_TorchPickle(): opts = mt.MapOptions() tmap = mt.CreateTriangular(dim,dim,3,opts) # Simple third order map - tmap2 = tmap.torch(store_coeffs=False) - x = torch.randn(numSamps, dim, dtype=torch.double) - coeffs = torch.randn(tmap.numCoeffs, dtype=torch.double) + tmap2 = tmap.torch(store_coeffs=True) + y = tmap2.forward(x) + - y = tmap2(x,coeffs) - assert y.shape[0] == numSamps - assert y.shape[1] == dim - assert not y.isnan().any() + map_bytes = dill.dumps(tmap2, dill.HIGHEST_PROTOCOL) + tmap3 = dill.loads(map_bytes) + + y2 = tmap3.forward(x) + assert (y2-y).abs().max() < 1e-8 if __name__=='__main__': @@ -118,4 +134,5 @@ def test_AutogradCoeffAsInput(): test_Autograd() test_AutogradCoeffs() test_TorchMethod() + test_TorchPickle() test_AutogradCoeffAsInput() \ No newline at end of file From d035eee2879565a6457c736198ab1720a48c5fca Mon Sep 17 00:00:00 2001 From: Matthew Parno Date: Tue, 27 Feb 2024 09:13:59 -0500 Subject: [PATCH 2/6] Bumped version and added dill dependency so tests can run. --- .docker/environment.yml | 1 + .github/environment.yml | 2 +- bindings/python/package/torch.py | 3 --- pyproject.toml | 2 +- 4 files changed, 3 insertions(+), 5 deletions(-) diff --git a/.docker/environment.yml b/.docker/environment.yml index 9858ab5e..e470641a 100644 --- a/.docker/environment.yml +++ b/.docker/environment.yml @@ -12,3 +12,4 @@ dependencies: - gcc - gxx - make + - dill diff --git a/.github/environment.yml b/.github/environment.yml index 4864ac20..5bf2b14c 100644 --- a/.github/environment.yml +++ b/.github/environment.yml @@ -13,4 +13,4 @@ dependencies: - cereal >= 1.3 - nlopt >= 2.7 - pytorch - + - dill diff --git a/bindings/python/package/torch.py b/bindings/python/package/torch.py index d2a1334e..6882c618 100644 --- a/bindings/python/package/torch.py +++ b/bindings/python/package/torch.py @@ -1,10 +1,7 @@ import torch -import dill - from .torch_helpers import ExtractTorchTensorData, MpartTorchAutograd - class TorchParameterizedFunctionBase(torch.nn.Module): """ Defines a wrapper around the MParT ParameterizedFunctionBase class that can be used with pytorch. diff --git a/pyproject.toml b/pyproject.toml index 123404c0..56a8c3ad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ license={file="LICENSE.txt"} readme="README.md" requires-python = ">=3.7" description="A Monotone Parameterization Toolkit" -version="2.2.1" +version="2.2.2" keywords=["Measure Transport", "Monotone", "Transport Map", "Isotonic Regression", "Triangular", "Knothe-Rosenblatt"] [project.urls] From 0b5e6e39e260f45a0a647e4ebafd53b1f5d4e68f Mon Sep 17 00:00:00 2001 From: Matthew Parno Date: Tue, 27 Feb 2024 09:16:38 -0500 Subject: [PATCH 3/6] Added release branch to github workflows. --- .github/workflows/build-bindings.yml | 1 + .github/workflows/build-doc.yml | 1 + .github/workflows/build-external-lib-tests.yml | 1 + .github/workflows/build-push-docker.yml | 1 + .github/workflows/build-tests.yml | 1 + 5 files changed, 5 insertions(+) diff --git a/.github/workflows/build-bindings.yml b/.github/workflows/build-bindings.yml index 0f029c4d..4856f39c 100644 --- a/.github/workflows/build-bindings.yml +++ b/.github/workflows/build-bindings.yml @@ -3,6 +3,7 @@ name: binding-tests on: push: branches: + - release - main pull_request: {} diff --git a/.github/workflows/build-doc.yml b/.github/workflows/build-doc.yml index 1214797c..17da50c5 100644 --- a/.github/workflows/build-doc.yml +++ b/.github/workflows/build-doc.yml @@ -4,6 +4,7 @@ on: push: branches: - main + - release jobs: build-docs: diff --git a/.github/workflows/build-external-lib-tests.yml b/.github/workflows/build-external-lib-tests.yml index c8cbb73c..bd13ae51 100644 --- a/.github/workflows/build-external-lib-tests.yml +++ b/.github/workflows/build-external-lib-tests.yml @@ -4,6 +4,7 @@ on: push: branches: - main + - release pull_request: {} env: diff --git a/.github/workflows/build-push-docker.yml b/.github/workflows/build-push-docker.yml index 2c784397..336a28b7 100644 --- a/.github/workflows/build-push-docker.yml +++ b/.github/workflows/build-push-docker.yml @@ -4,6 +4,7 @@ on: push: branches: - main + - release jobs: docker: diff --git a/.github/workflows/build-tests.yml b/.github/workflows/build-tests.yml index 8fabb02a..2ba6d1bf 100644 --- a/.github/workflows/build-tests.yml +++ b/.github/workflows/build-tests.yml @@ -4,6 +4,7 @@ on: push: branches: - main + - release pull_request: {} jobs: From c93aa58c762dd00ef0a7deb4c50074289e323aae Mon Sep 17 00:00:00 2001 From: Matthew Parno Date: Tue, 27 Feb 2024 09:36:54 -0500 Subject: [PATCH 4/6] Added missing include. --- tests/Distributions/Test_SampleGenerator.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/Distributions/Test_SampleGenerator.cpp b/tests/Distributions/Test_SampleGenerator.cpp index 36931920..58abf895 100644 --- a/tests/Distributions/Test_SampleGenerator.cpp +++ b/tests/Distributions/Test_SampleGenerator.cpp @@ -1,5 +1,8 @@ #include "Test_Distributions_Common.h" +#include +#include + TEST_CASE( "Testing SampleGenerator", "[SampleGenerator]") { // Sample 1000 points // Check empirical CDF against uniform CDF From 66dbc3dcd7c3d95853aa91c5eaace3a5da3fe3bd Mon Sep 17 00:00:00 2001 From: Daniel <43151183+dannys4@users.noreply.github.com> Date: Mon, 1 Apr 2024 17:10:20 -0400 Subject: [PATCH 5/6] Fix GitHub CI on release (#405) --- .github/workflows/build-bindings.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/build-bindings.yml b/.github/workflows/build-bindings.yml index 4856f39c..1ab32340 100644 --- a/.github/workflows/build-bindings.yml +++ b/.github/workflows/build-bindings.yml @@ -48,7 +48,7 @@ jobs: - name: Setup Julia run: | - julia -e "using Pkg; Pkg.add([\"CxxWrap\",\"TestReports\"])" + julia -e "using Pkg; Pkg.add([Pkg.PackageSpec(;name=\"CxxWrap\",version=v\"0.14.2\"),Pkg.PackageSpec(;name=\"TestReports\")])" export GITHUB_JULIA_PATH=$(julia -e "println(DEPOT_PATH[1])") echo -n $'[bee5971c-294f-5168-9fcd-9fb3c811d495]\nMParT = \"' >> $GITHUB_JULIA_PATH/artifacts/Overrides.toml echo -n $GITHUB_WORKSPACE >> $GITHUB_JULIA_PATH/artifacts/Overrides.toml From 1ae526b6c53a3a70d9ee9e9bd636afd4e3fca584 Mon Sep 17 00:00:00 2001 From: Daniel <43151183+dannys4@users.noreply.github.com> Date: Tue, 2 Apr 2024 07:02:46 -0400 Subject: [PATCH 6/6] Fix issue 286 (#403) * Fix issue 286 * Add binding change --- bindings/python/src/MultiIndex.cpp | 32 +++++++++++++++++++++ bindings/python/tests/test_MultiIndexSet.py | 8 ++++++ 2 files changed, 40 insertions(+) diff --git a/bindings/python/src/MultiIndex.cpp b/bindings/python/src/MultiIndex.cpp index d97c9614..2e01bcd0 100644 --- a/bindings/python/src/MultiIndex.cpp +++ b/bindings/python/src/MultiIndex.cpp @@ -24,6 +24,37 @@ namespace py = pybind11; using namespace mpart::binding; +template +using Matrix_Map_T = Eigen::Map, 0, Eigen::Stride>; + +mpart::MultiIndexSet MultiIndexSet_PyBuffer(py::buffer x){ + constexpr bool rowMajor = Eigen::Matrix::Flags & Eigen::RowMajorBit; + + py::buffer_info info = x.request(); + + // Check for int32, int64 + bool is_int32 = info.format == py::format_descriptor::format(); + bool is_int64 = info.format == "l"; // This is based on a pybind bug; numpy int64 buffer is l, not q + if (!(is_int32 || is_int64)) + throw std::runtime_error("Incompatible format: expected an array of either int32 or int64!"); + + if (info.ndim != 2) + throw std::runtime_error("Expected array with ndims = 2"); + + int stride_size = is_int32 ? sizeof(int32_t) : sizeof(int64_t); + Eigen::Stride strides( + info.strides[rowMajor ? 0 : 1] / (py::ssize_t)stride_size, + info.strides[rowMajor ? 1 : 0] / (py::ssize_t)stride_size + ); + + if(is_int64) { // Is int64 + Matrix_Map_T map_64 (static_cast(info.ptr), info.shape[0], info.shape[1], strides); + return mpart::MultiIndexSet {map_64.cast()}; + } else { // Is int32 + Matrix_Map_T map (static_cast(info.ptr), info.shape[0], info.shape[1], strides); + return mpart::MultiIndexSet {map}; + } +} void mpart::binding::MultiIndexWrapper(py::module &m) { @@ -111,6 +142,7 @@ void mpart::binding::MultiIndexWrapper(py::module &m) // MultiIndexSet py::class_>(m, "MultiIndexSet") .def(py::init()) + .def(py::init>(&MultiIndexSet_PyBuffer)) .def(py::init const&>()) .def("fix", &MultiIndexSet::Fix) .def("__len__", &MultiIndexSet::Length, "Retrieves the length of _each_ multiindex within this set (i.e. the dimension of the input)") diff --git a/bindings/python/tests/test_MultiIndexSet.py b/bindings/python/tests/test_MultiIndexSet.py index 9b70ec4f..85fa4f3d 100644 --- a/bindings/python/tests/test_MultiIndexSet.py +++ b/bindings/python/tests/test_MultiIndexSet.py @@ -10,6 +10,14 @@ msetTensorProduct = mpart.MultiIndexSet.CreateTensorProduct(3,4,noneLim) msetTotalOrder = mpart.MultiIndexSet.CreateTotalOrder(3,4,noneLim) +def test_create(): + mset_one = mpart.MultiIndexSet([[2]]) + assert mset_one.Size() == 1 + assert len(mset_one) == 1 + mset_one = mpart.MultiIndexSet(np.array([[2]])) + assert mset_one.Size() == 1 + assert len(mset_one) == 1 + def test_max_degrees(): assert np.all(msetFromArray.MaxOrders() == [2,1])