LabeledTensor.h

// Copyright 2023, UChicago Argonne, LLC
// All Rights Reserved
// Software Name: NEML2 -- the New Engineering material Model Library, version 2
// By: Argonne National Laboratory
// OPEN SOURCE LICENSE (MIT)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
 
#pragma once
 
#include "neml2/misc/types.h"
#include "neml2/tensors/LabeledAxis.h"
#include "neml2/tensors/BatchTensor.h"
 
namespace neml2
{
template <class Derived, TorchSize D>
class LabeledTensor
{
public:
  LabeledTensor() = default;
 
  LabeledTensor(const torch::Tensor & tensor,
                TorchSize batch_dim,
                const std::vector<const LabeledAxis *> & axes);
 
  LabeledTensor(const BatchTensor & tensor, const std::vector<const LabeledAxis *> & axes);
 
  LabeledTensor(const Derived & other);
 
  void operator=(const Derived & other);
 
  // Should we mark it explicit?
  operator BatchTensor() const;
 
  // Should we mark it explicit?
  operator torch::Tensor() const;
 
  [[nodiscard]] static Derived
  empty(TorchShapeRef batch_shape,
        const std::vector<const LabeledAxis *> & axes,
        const torch::TensorOptions & options = default_tensor_options());
 
  [[nodiscard]] static Derived empty_like(const Derived & other);
 
  [[nodiscard]] static Derived
  zeros(TorchShapeRef batch_shape,
        const std::vector<const LabeledAxis *> & axes,
        const torch::TensorOptions & options = default_tensor_options());
 
  [[nodiscard]] static Derived zeros_like(const Derived & other);
 
  Derived clone(torch::MemoryFormat memory_format = torch::MemoryFormat::Contiguous) const;
 
  template <typename T>
  void copy_(const T & other);
 
  Derived detach() const;
 
  void detach_();
 
  void zero_();
 
  const BatchTensor & tensor() const { return _tensor; }
  BatchTensor & tensor() { return _tensor; }
 
  torch::TensorOptions options() const { return _tensor.options(); }
 
  TorchSize batch_dim() const;
 
  TorchSize base_dim() const;
 
  TorchShapeRef batch_sizes() const;
 
  TorchShapeRef base_sizes() const;
 
  const std::vector<const LabeledAxis *> & axes() const { return _axes; }
 
  const LabeledAxis & axis(TorchSize i = 0) const { return *_axes[i]; }
 
  template <typename... S>
  TorchSlice slice_indices(S &&... names) const;
 
  TorchShapeRef storage_size() const;
 
  template <typename... S>
  TorchShape storage_size(S &&... names) const;
 
  template <typename... S>
  BatchTensor operator()(S &&... names) const;
 
  Derived slice(TorchSize i, const std::string & name) const;
 
  template <typename... S>
  Derived block(S &&... names) const;
 
  Derived batch_index(TorchSlice indices) const;
 
  void batch_index_put(TorchSlice indices, const torch::Tensor & other);
 
  BatchTensor base_index(TorchSlice indices) const;
 
  void base_index_put(TorchSlice indices, const torch::Tensor & other);
 
  template <typename T>
  struct variable_type
  {
    typedef T type;
  };
 
  template <typename T, typename... S>
  typename variable_type<T>::type get(S &&... names) const
  {
    return T((*this)(names...).view(utils::add_shapes(batch_sizes(), T::const_base_sizes)),
             batch_dim());
  }
 
  template <typename T, typename... S>
  typename variable_type<T>::type get_list(S &&... names) const
  {
    return T(((*this)(names...))
                 .reshape(utils::add_shapes(this->batch_sizes(), -1, T::const_base_sizes)),
             this->batch_dim() + sizeof...(names));
  }
 
  template <typename T, typename... S>
  void set(const BatchTensorBase<T> & value, S &&... names)
  {
    (*this)(names...).index_put_(
        {torch::indexing::None},
        value.reshape(utils::add_shapes(value.batch_sizes(), storage_size(names...))));
  }
 
  template <typename T, typename... S>
  void set_list(const BatchTensorBase<T> & value, S &&... names)
  {
    this->set(BatchTensor(value, value.batch_dim() - sizeof...(names)), names...);
  }
 
  Derived operator-() const;
 
  Derived to(const torch::TensorOptions & options) const;
 
protected:
  BatchTensor _tensor;
 
  // Urgh, I can't use const references here as the elements of a vector has to be "assignable".
  std::vector<const LabeledAxis *> _axes;
 
private:
  template <std::size_t... I, typename... S>
  TorchSlice slice_indices_impl(std::index_sequence<I...>, S &&... names) const;
 
  template <std::size_t... I, typename... S>
  TorchShape storage_size_impl(std::index_sequence<I...>, S &&... names) const;
 
  template <std::size_t... I, typename... S>
  Derived block_impl(std::index_sequence<I...>, S &&... names) const;
};
 
template <class Derived, TorchSize D>
template <typename T>
void
LabeledTensor<Derived, D>::copy_(const T & other)
{
  _tensor.copy_(other);
}
 
template <class Derived, TorchSize D>
template <typename... S>
TorchSlice
LabeledTensor<Derived, D>::slice_indices(S &&... names) const
{
  static_assert(sizeof...(names) == D, "Wrong labaled dimesion in LabeledTensor::slice_indices");
  return slice_indices_impl(std::make_index_sequence<sizeof...(names)>(),
                            std::forward<S>(names)...);
}
 
template <class Derived, TorchSize D>
template <std::size_t... I, typename... S>
TorchSlice
LabeledTensor<Derived, D>::slice_indices_impl(std::index_sequence<I...>, S &&... names) const
{
  return {_axes[I]->indices(names)...};
}
 
template <class Derived, TorchSize D>
template <typename... S>
TorchShape
LabeledTensor<Derived, D>::storage_size(S &&... names) const
{
  static_assert(sizeof...(names) == D, "Wrong labaled dimesion in LabeledTensor::storage_size");
  return storage_size_impl(std::make_index_sequence<D>(), std::forward<S>(names)...);
}
 
template <class Derived, TorchSize D>
template <std::size_t... I, typename... S>
TorchShape
LabeledTensor<Derived, D>::storage_size_impl(std::index_sequence<I...>, S &&... names) const
{
  return {_axes[I]->storage_size(names)...};
}
 
template <class Derived, TorchSize D>
template <typename... S>
BatchTensor
LabeledTensor<Derived, D>::operator()(S &&... names) const
{
  static_assert(sizeof...(names) == D, "Wrong labeled dimension in LabeledTensor::operator()");
  return base_index(slice_indices(names...));
}
 
template <class Derived, TorchSize D>
template <typename... S>
Derived
LabeledTensor<Derived, D>::block(S &&... names) const
{
  return block_impl(std::make_index_sequence<sizeof...(names)>(), std::forward<S>(names)...);
}
 
template <class Derived, TorchSize D>
template <std::size_t... I, typename... S>
Derived
LabeledTensor<Derived, D>::block_impl(std::index_sequence<I...>, S &&... names) const
{
  TorchSlice idx = {_axes[I]->indices(names)...};
  std::vector<const LabeledAxis *> new_axes = {&_axes[I]->subaxis(names)...};
  return Derived(_tensor.base_index(idx), new_axes);
}
} // namespace neml2