{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE NoStarIsType #-}

-- | __Tensor networks__ as a hypergraph category: the same arrows as the matrices of
-- "Proarrow.Category.Instance.Mat", kept as a network instead of as one matrix. An object is a list
-- of dimensions, one for each wire. An arrow has nodes, each with a dimension, a node for each of its
-- input and output wires, and dense factors: flat vectors, each with a node for each of its axes. Its
-- entry at given indices of the wires is the sum, over the indices of the nodes that agree with
-- those of the wires, of the products of the entries of the factors.
--
-- So the structure of a hypergraph category costs nothing: identities, swaps, copying, merging,
-- discarding, cups and caps are wirings without factors, and the tensor puts factors side by side
-- without multiplying them. Composition glues the wirings and sums out the nodes that are no longer
-- on a wire, multiplying only the factors that a summed node joins. With "Proarrow.Tools.Einsum" a
-- network of tensors is contracted in the order its read-back chooses.
--
-- With the package's @blas@ flag, a contraction of two factors that is a matrix product is handed
-- to the system's BLAS (Accelerate on macOS, OpenBLAS elsewhere) at entries of type 'P.Double',
-- 'P.Float' and @'Complex' 'P.Double'@.
--
-- The biproducts are direct sums, a single wire whose dimension is the sum of the sizes of the two
-- objects. Their injections, projections and pairings are dense matrices, and so are the
-- distributors.
module Proarrow.Category.Instance.TensorNetwork
  ( TNET (..)
  , TensorNetwork
  , Scalar
  , dimsOf
  , Size
  , fromVector
  , toVector
  , fromEntries
  , entries
  , matMul
  ) where

import Data.Complex (Complex, conjugate)
import Data.Containers.ListUtils (nubOrd)
import Data.IntMap.Strict qualified as IM
import Data.IntSet qualified as IS
import Data.Kind (Constraint, Type)
import Data.List qualified as List
import Data.Proxy (Proxy (..))
import Data.Set qualified as Set
import Data.Type.Equality ((:~:) (..))
import Data.Vector.Storable qualified as SV
import Data.Vector.Storable.Mutable qualified as MSV
import Data.Vector.Unboxed qualified as UV
import Foreign.Storable (Storable)
import GHC.TypeNats (KnownNat, Nat, SNat, natVal, sameNat, withKnownNat, withSomeSNat, type (*), type (+))
import Unsafe.Coerce (unsafeCoerce)
import Prelude (Int, ($), (*), (+), (-), (==))
import Prelude qualified as P

import Proarrow.Category.Enriched.Dagger (DaggerProfunctor (..))
import Proarrow.Category.Instance.FinHask (unionFind)
import Proarrow.Category.Instance.TensorNetwork.Blas (Gemm, gemmComplexDouble, gemmDouble, gemmFloat)
import Proarrow.Category.Monoidal (Monoidal (..), MonoidalProfunctor (..), SymMonoidal (..))
import Proarrow.Category.Monoidal.Action (MonoidalAction)
import Proarrow.Category.Monoidal.Closed (Closed (..))
import Proarrow.Category.Monoidal.CompactClosed (CompactClosed (..), coactCC)
import Proarrow.Category.Monoidal.CopyDiscard (CopyDiscard)
import Proarrow.Category.Monoidal.Dialogue (Dialogue (..))
import Proarrow.Category.Monoidal.Distributive (Distributive (..))
import Proarrow.Category.Monoidal.Hypergraph (Frobenius, Hypergraph, Sized (..), cap, cup)
import Proarrow.Category.Monoidal.IsoMix (IsoMix (..))
import Proarrow.Category.Monoidal.StarAutonomous (ExpSA, StarAutonomous (..), applySA, currySA, expSA)
import Proarrow.Category.Monoidal.Strength (Costrong (..))
import Proarrow.Colimit.BinaryCoproduct (HasBinaryCoproducts (..), HasBiproducts)
import Proarrow.Colimit.Initial (HasInitialObject (..))
import Proarrow.Core (CAT, CategoryOf (..), Is, Profunctor (..), Promonad (..), UN, dimapDefault, type (+->))
import Proarrow.Limit.BinaryProduct (HasBinaryProducts (..))
import Proarrow.Limit.Terminal (HasTerminalObject (..))
import Proarrow.Monoid (CocommutativeComonoid, CommutativeMonoid, Comonoid (..), Monoid (..))
import Proarrow.Object (KnownListOf (..), eqListOf, mapListOf, withKnownAppend, withKnownAssoc, type (++))
import Proarrow.Optic.Iso (DecidableIso (..), isoFromEquality)
import Proarrow.Tools.ContractionOrder qualified as Plan

-- | The entries of the factors: numbers that can be stored unboxed. An instance needs no methods:
-- the loop that contracts factors is then compiled for its type.
type Scalar :: Type -> Constraint
class (P.Num e, Storable e) => Scalar e where
  -- | The entries of a contraction of two factors: given the dimension of each axis of the result
  -- and how far a step along it moves in the entries of each factor, the offset into each factor
  -- of each summed index, and the entries of the two factors, each entry is a sum of products.
  kernel :: [(Dim, Int, Int)] -> UV.Vector Index -> UV.Vector Index -> SV.Vector e -> SV.Vector e -> SV.Vector e
  -- inlined into each instance, so that the loop is specialised to its type
  kernel = [(Node, Node, Node)]
-> Vector Node -> Vector Node -> Vector e -> Vector e -> Vector e
forall e.
(Num e, Storable e) =>
[(Node, Node, Node)]
-> Vector Node -> Vector Node -> Vector e -> Vector e -> Vector e
contractLoop
  {-# INLINE kernel #-}

  -- | The conjugation of entries, when it is not the identity, as for complex numbers.
  conj :: P.Maybe (e -> e)
  conj = Maybe (e -> e)
forall a. Maybe a
P.Nothing

  -- | A fast matrix product, which a contraction of two factors that is one is handed to.
  gemm :: P.Maybe (Gemm e)
  gemm = Maybe (Gemm e)
forall a. Maybe a
P.Nothing

instance Scalar Int
instance Scalar P.Double where
  gemm :: Maybe (Gemm Double)
gemm = Maybe (Gemm Double)
gemmDouble
instance Scalar P.Float where
  gemm :: Maybe (Gemm Float)
gemm = Maybe (Gemm Float)
gemmFloat
instance Scalar (Complex P.Double) where
  conj :: Maybe (Complex Double -> Complex Double)
conj = (Complex Double -> Complex Double)
-> Maybe (Complex Double -> Complex Double)
forall a. a -> Maybe a
P.Just Complex Double -> Complex Double
forall a. Num a => Complex a -> Complex a
conjugate
  gemm :: Maybe (Gemm (Complex Double))
gemm = Maybe (Gemm (Complex Double))
gemmComplexDouble

-- | The loop that 'kernel' runs, written once for every type: a loop over each axis of the result,
-- the last one innermost, reading the entries of both factors with their steps, and for each
-- entry a loop over the summed indices.
{-# INLINEABLE contractLoop #-}
contractLoop
  :: (P.Num e, Storable e)
  => [(Dim, Int, Int)] -> UV.Vector Index -> UV.Vector Index -> SV.Vector e -> SV.Vector e -> SV.Vector e
contractLoop :: forall e.
(Num e, Storable e) =>
[(Node, Node, Node)]
-> Vector Node -> Vector Node -> Vector e -> Vector e -> Vector e
contractLoop [(Node, Node, Node)]
steps !Vector Node
so !Vector Node
to !Vector e
v !Vector e
w = (forall s. ST s (MVector s e)) -> Vector e
forall a. Storable a => (forall s. ST s (MVector s a)) -> Vector a
SV.create do
  out <- Node -> ST s (MVector (PrimState (ST s)) e)
forall (m :: Type -> Type) a.
(PrimMonad m, Storable a) =>
Node -> m (MVector (PrimState m) a)
MSV.unsafeNew ([Node] -> Node
forall a. Num a => [a] -> a
forall (t :: Type -> Type) a. (Foldable t, Num a) => t a -> a
P.product [Node
d | (Node
d, Node
_, Node
_) <- [(Node, Node, Node)]
steps])
  let go [] !Node
a !Node
b !Node
dst = MVector (PrimState (ST s)) e -> Node -> e -> ST s ()
forall (m :: Type -> Type) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Node -> a -> m ()
MSV.unsafeWrite MVector s e
MVector (PrimState (ST s)) e
out Node
dst (Node -> Node -> e
entry Node
a Node
b)
      go [(!Node
d, !Node
s, !Node
t)] !Node
a !Node
b !Node
dst
        | Node
n Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
== Node
1 = Node -> Node -> Node -> ST s ()
products Node
0 (Node
a Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Vector Node
so Vector Node -> Node -> Node
forall a. Unbox a => Vector a -> Node -> a
`UV.unsafeIndex` Node
0) (Node
b Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Vector Node
to Vector Node -> Node -> Node
forall a. Unbox a => Vector a -> Node -> a
`UV.unsafeIndex` Node
0)
        | Bool
P.otherwise = Node -> Node -> Node -> ST s ()
sums Node
0 Node
a Node
b
        where
          -- with nothing summed, each entry is one product
          products :: Node -> Node -> Node -> ST s ()
products !Node
i !Node
x !Node
y
            | Node
i Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
== Node
d = () -> ST s ()
forall a. a -> ST s a
forall (f :: Type -> Type) a. Applicative f => a -> f a
P.pure ()
            | Bool
P.otherwise =
                MVector (PrimState (ST s)) e -> Node -> e -> ST s ()
forall (m :: Type -> Type) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Node -> a -> m ()
MSV.unsafeWrite MVector s e
MVector (PrimState (ST s)) e
out (Node
dst Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
i) (Vector e
v Vector e -> Node -> e
forall a. Storable a => Vector a -> Node -> a
`SV.unsafeIndex` Node
x e -> e -> e
forall a. Num a => a -> a -> a
* Vector e
w Vector e -> Node -> e
forall a. Storable a => Vector a -> Node -> a
`SV.unsafeIndex` Node
y) ST s () -> ST s () -> ST s ()
forall a b. ST s a -> ST s b -> ST s b
forall (m :: Type -> Type) a b. Monad m => m a -> m b -> m b
P.>> Node -> Node -> Node -> ST s ()
products (Node
i Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
1) (Node
x Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
s) (Node
y Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
t)
          sums :: Node -> Node -> Node -> ST s ()
sums !Node
i !Node
x !Node
y
            | Node
i Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
== Node
d = () -> ST s ()
forall a. a -> ST s a
forall (f :: Type -> Type) a. Applicative f => a -> f a
P.pure ()
            | Bool
P.otherwise = MVector (PrimState (ST s)) e -> Node -> e -> ST s ()
forall (m :: Type -> Type) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Node -> a -> m ()
MSV.unsafeWrite MVector s e
MVector (PrimState (ST s)) e
out (Node
dst Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
i) (Node -> Node -> e
entry Node
x Node
y) ST s () -> ST s () -> ST s ()
forall a b. ST s a -> ST s b -> ST s b
forall (m :: Type -> Type) a b. Monad m => m a -> m b -> m b
P.>> Node -> Node -> Node -> ST s ()
sums (Node
i Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
1) (Node
x Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
s) (Node
y Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
t)
      go ((!Node
d, !Node
s, !Node
t) : [(Node, Node, Node)]
rest) !Node
a !Node
b !Node
dst = Node -> Node -> Node -> ST s ()
outer Node
0 Node
a Node
b
        where
          size' :: Node
size' = [Node] -> Node
forall a. Num a => [a] -> a
forall (t :: Type -> Type) a. (Foldable t, Num a) => t a -> a
P.product [Node
e | (Node
e, Node
_, Node
_) <- [(Node, Node, Node)]
rest]
          outer :: Node -> Node -> Node -> ST s ()
outer !Node
i !Node
x !Node
y
            | Node
i Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
== Node
d = () -> ST s ()
forall a. a -> ST s a
forall (f :: Type -> Type) a. Applicative f => a -> f a
P.pure ()
            | Bool
P.otherwise = [(Node, Node, Node)] -> Node -> Node -> Node -> ST s ()
go [(Node, Node, Node)]
rest Node
x Node
y (Node
dst Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
i Node -> Node -> Node
forall a. Num a => a -> a -> a
* Node
size') ST s () -> ST s () -> ST s ()
forall a b. ST s a -> ST s b -> ST s b
forall (m :: Type -> Type) a b. Monad m => m a -> m b -> m b
P.>> Node -> Node -> Node -> ST s ()
outer (Node
i Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
1) (Node
x Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
s) (Node
y Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
t)
  go steps 0 0 0
  P.pure out
  where
    n :: Node
n = Vector Node -> Node
forall a. Unbox a => Vector a -> Node
UV.length Vector Node
so
    -- the sum of the products at a place in each factor; its arguments are strict, so that the loop
    -- does not look them up again for every index
    entry :: Node -> Node -> e
entry !Node
a !Node
b = Node -> e -> e
sumFrom Node
0 e
0
      where
        sumFrom :: Node -> e -> e
sumFrom !Node
i !e
acc =
          if Node
i Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
== Node
n
            then e
acc
            else
              Node -> e -> e
sumFrom (Node
i Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
1) (e
acc e -> e -> e
forall a. Num a => a -> a -> a
+ Vector e
v Vector e -> Node -> e
forall a. Storable a => Vector a -> Node -> a
`SV.unsafeIndex` (Node
a Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Vector Node
so Vector Node -> Node -> Node
forall a. Unbox a => Vector a -> Node -> a
`UV.unsafeIndex` Node
i) e -> e -> e
forall a. Num a => a -> a -> a
* Vector e
w Vector e -> Node -> e
forall a. Storable a => Vector a -> Node -> a
`SV.unsafeIndex` (Node
b Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Vector Node
to Vector Node -> Node -> Node
forall a. Unbox a => Vector a -> Node -> a
`UV.unsafeIndex` Node
i))

-- | A node of a network, numbered from 0.
type Node = Int

-- | The dimension of a wire or a node.
type Dim = Int

-- | An index along a wire or an axis, or a position in a factor's entries.
type Index = Int

-- | Objects are lists of dimensions, one for each wire.
type data TNET (e :: Type) = TN [Nat]

-- | The dimensions of a list, as numbers.
dimsOf :: forall ns. (KnownListOf KnownNat ns) => [Dim]
dimsOf :: forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf = forall {k} (c :: OB k) (xs :: [k]) r.
(forall (x :: k). c x => r) -> ListOf c xs -> [r]
forall (c :: OB Nat) (xs :: [Nat]) r.
(forall (x :: Nat). c x => r) -> ListOf c xs -> [r]
mapListOf @KnownNat (\ @n -> Nat -> Node
forall a b. (Integral a, Num b) => a -> b
P.fromIntegral (Proxy x -> Nat
forall (n :: Nat) (proxy :: Nat -> Type).
KnownNat n =>
proxy n -> Nat
natVal (forall (t :: Nat). Proxy t
forall {k} (t :: k). Proxy t
Proxy @n))) (forall {k} (c :: OB k) (xs :: [k]). KnownListOf c xs => ListOf c xs
forall (c :: OB Nat) (xs :: [Nat]). KnownListOf c xs => ListOf c xs
listOf @KnownNat @ns)

-- | A dense factor: its entries, row by row over its axes, the node of each axis, and the
-- conjugates of its entries. Both are computed when first needed, so the dagger can swap them, and
-- each factor conjugates its entries at most once.
data Factor e = Factor {forall e. Factor e -> [Node]
axes :: [Node], forall e. Factor e -> Vector e
values :: ~(SV.Vector e), forall e. Factor e -> Vector e
conjugates :: ~(SV.Vector e)}

-- | The factor with the given axes and entries.
factor :: (Scalar e) => [Node] -> SV.Vector e -> Factor e
factor :: forall e. Scalar e => [Node] -> Vector e -> Factor e
factor [Node]
as Vector e
vs = [Node] -> Vector e -> Vector e -> Factor e
forall e. [Node] -> Vector e -> Vector e -> Factor e
Factor [Node]
as Vector e
vs (Vector e -> ((e -> e) -> Vector e) -> Maybe (e -> e) -> Vector e
forall b a. b -> (a -> b) -> Maybe a -> b
P.maybe Vector e
vs ((e -> e) -> Vector e -> Vector e
forall a b.
(Storable a, Storable b) =>
(a -> b) -> Vector a -> Vector b
`SV.map` Vector e
vs) Maybe (e -> e)
forall e. Scalar e => Maybe (e -> e)
conj)

-- | The nodes, numbered from 0 with their dimensions, the node of each input and output wire, and
-- the factors. Kept so that every node is on a wire and no factor has an axis twice.
data Net e = Net {forall e. Net e -> [Node]
nodeDims :: [Dim], forall e. Net e -> [Node]
ins :: [Node], forall e. Net e -> [Node]
outs :: [Node], forall e. Net e -> [Factor e]
factors :: [Factor e]}

-- | An arrow between two lists of dimensions.
type TensorNetwork :: CAT (TNET e)
data TensorNetwork a b where
  TensorNetwork
    :: forall {e} as bs. (KnownListOf KnownNat as, KnownListOf KnownNat bs) => Net e -> TensorNetwork (TN as :: TNET e) (TN bs)

-- | A wiring without factors: nodes of the given dimensions, and the node of each input and output
-- wire.
wiring
  :: forall {e} as bs
   . (KnownListOf KnownNat as, KnownListOf KnownNat bs)
  => [Dim] -> [Node] -> [Node] -> TensorNetwork (TN as :: TNET e) (TN bs)
wiring :: forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
[Node] -> [Node] -> [Node] -> TensorNetwork (TN as) (TN bs)
wiring [Node]
dims [Node]
is [Node]
os = Net e -> TensorNetwork (TN as) (TN bs)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Net e -> TensorNetwork (TN as) (TN bs)
TensorNetwork ([Node] -> [Node] -> [Node] -> [Factor e] -> Net e
forall e. [Node] -> [Node] -> [Node] -> [Factor e] -> Net e
Net [Node]
dims [Node]
is [Node]
os [])

-- | The wires going straight through, between two lists of the same dimensions.
straight
  :: forall {e} as bs. (KnownListOf KnownNat as, KnownListOf KnownNat bs) => TensorNetwork (TN as :: TNET e) (TN bs)
straight :: forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
TensorNetwork (TN as) (TN bs)
straight = [Node] -> [Node] -> [Node] -> TensorNetwork (TN as) (TN bs)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
[Node] -> [Node] -> [Node] -> TensorNetwork (TN as) (TN bs)
wiring (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf @as) (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
wires @as) (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
wires @as)

-- | The arrow with the given entries, laid out as 'toVector' gives them. The length of the vector
-- must be the product of all the dimensions.
fromVector
  :: forall {e} as bs
   . (Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs)
  => SV.Vector e -> TensorNetwork (TN as :: TNET e) (TN bs)
fromVector :: forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Vector e -> TensorNetwork (TN as) (TN bs)
fromVector Vector e
vs
  | Vector e -> Node
forall a. Storable a => Vector a -> Node
SV.length Vector e
vs Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
P./= [Node] -> Node
forall a. Num a => [a] -> a
forall (t :: Type -> Type) a. (Foldable t, Num a) => t a -> a
P.product [Node]
dims = [Char] -> TensorNetwork (TN as) (TN bs)
forall a. HasCallStack => [Char] -> a
P.error [Char]
"fromVector: the length is not the product of the dimensions"
  | Bool
P.otherwise = Net e -> TensorNetwork (TN as) (TN bs)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Net e -> TensorNetwork (TN as) (TN bs)
TensorNetwork ([Node] -> [Node] -> [Node] -> [Factor e] -> Net e
forall e. [Node] -> [Node] -> [Node] -> [Factor e] -> Net e
Net [Node]
dims [Node]
is [Node]
os [[Node] -> Vector e -> Factor e
forall e. Scalar e => [Node] -> Vector e -> Factor e
factor ([Node]
os [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ [Node]
is) Vector e
vs])
  where
    dims :: [Node]
dims = forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf @as [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf @bs
    is :: [Node]
is = forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
wires @as
    os :: [Node]
os = [[Node] -> Node
forall a. [a] -> Node
forall (t :: Type -> Type) a. Foldable t => t a -> Node
P.length [Node]
is .. [Node] -> Node
forall a. [a] -> Node
forall (t :: Type -> Type) a. Foldable t => t a -> Node
P.length [Node]
dims Node -> Node -> Node
forall a. Num a => a -> a -> a
- Node
1]

-- | The entries, row by row: a row for each index of the output wires and a column for each index
-- of the input wires, both with the last wire varying fastest. This relies on every node of the
-- network being on a wire, which composition keeps so.
toVector :: forall {e} a b. (Scalar e) => TensorNetwork (a :: TNET e) b -> SV.Vector e
toVector :: forall {e} (a :: TNET e) (b :: TNET e).
Scalar e =>
TensorNetwork a b -> Vector e
toVector (TensorNetwork (Net [Node]
dims [Node]
is [Node]
os [Factor e]
fs))
  | [Node] -> Bool
distinct [Node]
ws = Factor e -> Vector e
forall e. Factor e -> Vector e
values ([Node] -> Factor e
product [Node]
ws)
  | Bool
P.otherwise = (forall s. ST s (MVector s e)) -> Vector e
forall a. Storable a => (forall s. ST s (MVector s a)) -> Vector a
SV.create do
      out <- Node -> e -> ST s (MVector (PrimState (ST s)) e)
forall (m :: Type -> Type) a.
(PrimMonad m, Storable a) =>
Node -> a -> m (MVector (PrimState m) a)
MSV.replicate ([Node] -> [Node] -> Node
extent [Node]
dims [Node]
ws) e
0
      let vs = Factor e -> Vector e
forall e. Factor e -> Vector e
values ([Node] -> Factor e
product [Node]
order)
          go [] !Node
src !Node
dst = MVector (PrimState (ST s)) e -> Node -> e -> ST s ()
forall (m :: Type -> Type) a.
(PrimMonad m, Storable a) =>
MVector (PrimState m) a -> Node -> a -> m ()
MSV.unsafeWrite MVector s e
MVector (PrimState (ST s)) e
out Node
dst (Vector e
vs Vector e -> Node -> e
forall a. Storable a => Vector a -> Node -> a
`SV.unsafeIndex` Node
src) ST s () -> ST s Node -> ST s Node
forall a b. ST s a -> ST s b -> ST s b
forall (m :: Type -> Type) a b. Monad m => m a -> m b -> m b
P.>> Node -> ST s Node
forall a. a -> ST s a
forall (f :: Type -> Type) a. Applicative f => a -> f a
P.pure (Node
src Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
1)
          go ((Node
d, Node
w) : [(Node, Node)]
rest) !Node
src !Node
dst = Node -> Node -> ST s Node
loop Node
0 Node
src
            where
              loop :: Node -> Node -> ST s Node
loop !Node
i !Node
from
                | Node
i Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
== Node
d = Node -> ST s Node
forall a. a -> ST s a
forall (f :: Type -> Type) a. Applicative f => a -> f a
P.pure Node
from
                | Bool
P.otherwise = [(Node, Node)] -> Node -> Node -> ST s Node
go [(Node, Node)]
rest Node
from (Node
dst Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
i Node -> Node -> Node
forall a. Num a => a -> a -> a
* Node
w) ST s Node -> (Node -> ST s Node) -> ST s Node
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
P.>>= Node -> Node -> ST s Node
loop (Node
i Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
1)
      _ <- go [(dims P.!! n, P.sum [st | (x, st) <- P.zip ws wireStrides, x == n]) | n <- order] 0 0
      P.pure out
  where
    ws :: [Node]
ws = [Node]
os [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ [Node]
is
    -- the product of the factors, kept over the given nodes
    product :: [Node] -> Factor e
product [Node]
keep = [Node] -> Factor e -> [Node] -> [Node] -> Factor e
forall e.
Scalar e =>
[Node] -> Factor e -> [Node] -> [Node] -> Factor e
contract1 [Node]
dims (case [Factor e]
fs of [] -> Factor e
forall e. Scalar e => Factor e
unit; Factor e
f : [Factor e]
rest -> (Factor e -> Factor e -> Factor e)
-> Factor e -> [Factor e] -> Factor e
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
P.foldl Factor e -> Factor e -> Factor e
times Factor e
f [Factor e]
rest) [Node]
keep []
    times :: Factor e -> Factor e -> Factor e
times Factor e
g Factor e
h = [Node] -> Factor e -> Factor e -> [Node] -> [Node] -> Factor e
forall e.
Scalar e =>
[Node] -> Factor e -> Factor e -> [Node] -> [Node] -> Factor e
contract [Node]
dims Factor e
g Factor e
h ([Node] -> [Node]
forall a. Ord a => [a] -> [a]
nubOrd (Factor e -> [Node]
forall e. Factor e -> [Node]
axes Factor e
g [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ Factor e -> [Node]
forall e. Factor e -> [Node]
axes Factor e
h)) []
    -- with a node on several wires, each entry of the product over the nodes goes to the place
    -- that its wires give it, and the others are 0
    order :: [Node]
order = [Node] -> [Node]
forall a. Ord a => [a] -> [a]
nubOrd [Node]
ws
    wireStrides :: [Node]
wireStrides = Node -> [Node] -> [Node]
forall a. Node -> [a] -> [a]
P.drop Node
1 ((Node -> Node -> Node) -> Node -> [Node] -> [Node]
forall a b. (a -> b -> b) -> b -> [a] -> [b]
P.scanr Node -> Node -> Node
forall a. Num a => a -> a -> a
(*) Node
1 ((Node -> Node) -> [Node] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap ([Node]
dims [Node] -> Node -> Node
forall a. HasCallStack => [a] -> Node -> a
P.!!) [Node]
ws))

-- | The arrow with the given rows of entries, laid out as 'entries' gives them.
fromEntries
  :: forall {e} as bs
   . (Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs)
  => [[e]] -> TensorNetwork (TN as :: TNET e) (TN bs)
fromEntries :: forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
[[e]] -> TensorNetwork (TN as) (TN bs)
fromEntries [[e]]
rows = Vector e -> TensorNetwork (TN as) (TN bs)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Vector e -> TensorNetwork (TN as) (TN bs)
fromVector ([e] -> Vector e
forall a. Storable a => [a] -> Vector a
SV.fromList ([[e]] -> [e]
forall (t :: Type -> Type) a. Foldable t => t [a] -> [a]
P.concat [[e]]
rows))

-- | The entries as rows, as 'toVector' lays them out.
entries :: forall {e} a b. (Scalar e) => TensorNetwork (a :: TNET e) b -> [[e]]
entries :: forall {e} (a :: TNET e) (b :: TNET e).
Scalar e =>
TensorNetwork a b -> [[e]]
entries f :: TensorNetwork a b
f@(TensorNetwork @as @bs Net e
_) = [Vector e -> [e]
forall a. Storable a => Vector a -> [a]
SV.toList (Node -> Node -> Vector e -> Vector e
forall a. Storable a => Node -> Node -> Vector a -> Vector a
SV.slice (Node
r Node -> Node -> Node
forall a. Num a => a -> a -> a
* Node
ni) Node
ni Vector e
v) | Node
r <- [Node
0 .. forall (ns :: [Nat]). KnownListOf KnownNat ns => Node
size @bs Node -> Node -> Node
forall a. Num a => a -> a -> a
- Node
1]]
  where
    v :: Vector e
v = TensorNetwork a b -> Vector e
forall {e} (a :: TNET e) (b :: TNET e).
Scalar e =>
TensorNetwork a b -> Vector e
toVector TensorNetwork a b
f
    ni :: Node
ni = forall (ns :: [Nat]). KnownListOf KnownNat ns => Node
size @as

-- | The matrix of a function from the indices of the inputs to those of the outputs: each column
-- has a 1 in the row the function gives it.
reindex
  :: forall {e} as bs
   . (Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs)
  => (Index -> Index) -> TensorNetwork (TN as :: TNET e) (TN bs)
reindex :: forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Node -> Node) -> TensorNetwork (TN as) (TN bs)
reindex Node -> Node
f = Vector e -> TensorNetwork (TN as) (TN bs)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Vector e -> TensorNetwork (TN as) (TN bs)
fromVector (Node -> (Node -> e) -> Vector e
forall a. Storable a => Node -> (Node -> a) -> Vector a
SV.generate (forall (ns :: [Nat]). KnownListOf KnownNat ns => Node
size @bs Node -> Node -> Node
forall a. Num a => a -> a -> a
* Node
ni) \Node
k -> let (Node
r, Node
c) = Node
k Node -> Node -> (Node, Node)
forall a. Integral a => a -> a -> (a, a)
`P.quotRem` Node
ni in if Node
r Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
== Node -> Node
f Node
c then e
1 else e
0)
  where
    ni :: Node
ni = forall (ns :: [Nat]). KnownListOf KnownNat ns => Node
size @as

-- | The arrow with no entries, into or out of a wire of dimension 0.
zero
  :: forall {e} as bs
   . (Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs)
  => TensorNetwork (TN as :: TNET e) (TN bs)
zero :: forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
TensorNetwork (TN as) (TN bs)
zero = Vector e -> TensorNetwork (TN as) (TN bs)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Vector e -> TensorNetwork (TN as) (TN bs)
fromVector Vector e
forall a. Storable a => Vector a
SV.empty

-- | The size of a list of dimensions: their product.
type Size :: [Nat] -> Nat
type family Size ns where
  Size '[] = 1
  Size (n ': ns) = n * Size ns

-- | The size of a list of dimensions, as a number.
size :: forall ns. (KnownListOf KnownNat ns) => Int
size :: forall (ns :: [Nat]). KnownListOf KnownNat ns => Node
size = [Node] -> Node
forall a. Num a => [a] -> a
forall (t :: Type -> Type) a. (Foldable t, Num a) => t a -> a
P.product (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf @ns)

-- | The dimension of the direct sum of two lists of dimensions.
withSum
  :: forall as bs r. (KnownListOf KnownNat as, KnownListOf KnownNat bs) => ((KnownNat (Size as + Size bs)) => r) -> r
withSum :: forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownNat (Size as + Size bs) => r) -> r
withSum = SNat (Size as + Size bs)
-> (KnownNat (Size as + Size bs) => r) -> r
forall (n :: Nat) r. SNat n -> (KnownNat n => r) -> r
withKnownNat SNat (Size as + Size bs)
n
  where
    n :: SNat (Size as + Size bs)
    n :: SNat (Size as + Size bs)
n = Nat
-> (forall (n :: Nat). SNat n -> SNat (Size as + Size bs))
-> SNat (Size as + Size bs)
forall r. Nat -> (forall (n :: Nat). SNat n -> r) -> r
withSomeSNat (Node -> Nat
forall a b. (Integral a, Num b) => a -> b
P.fromIntegral (forall (ns :: [Nat]). KnownListOf KnownNat ns => Node
size @as Node -> Node -> Node
forall a. Num a => a -> a -> a
+ forall (ns :: [Nat]). KnownListOf KnownNat ns => Node
size @bs)) SNat n -> SNat (Size as + Size bs)
forall (n :: Nat). SNat n -> SNat (Size as + Size bs)
forall a b. a -> b
unsafeCoerce

-- | The product of two factors, with the given nodes kept, in that order, and the others summed.
contract :: forall e. (Scalar e) => [Dim] -> Factor e -> Factor e -> [Node] -> [Node] -> Factor e
contract :: forall e.
Scalar e =>
[Node] -> Factor e -> Factor e -> [Node] -> [Node] -> Factor e
contract [Node]
dims Factor e
f Factor e
g [Node]
keep [Node]
summed
  | P.Just Gemm e
mm <- Maybe (Gemm e)
forall e. Scalar e => Maybe (Gemm e)
gemm, P.Just Factor e
r <- Gemm e
-> [Node]
-> Factor e
-> Factor e
-> [Node]
-> [Node]
-> Maybe (Factor e)
forall e.
Scalar e =>
Gemm e
-> [Node]
-> Factor e
-> Factor e
-> [Node]
-> [Node]
-> Maybe (Factor e)
viaGemm Gemm e
mm [Node]
dims Factor e
f Factor e
g [Node]
keep [Node]
summed = Factor e
r
  | Bool
P.otherwise =
      [Node] -> Vector e -> Factor e
forall e. Scalar e => [Node] -> Vector e -> Factor e
factor
        [Node]
keep
        ( [(Node, Node, Node)]
-> Vector Node -> Vector Node -> Vector e -> Vector e -> Vector e
forall e.
Scalar e =>
[(Node, Node, Node)]
-> Vector Node -> Vector Node -> Vector e -> Vector e -> Vector e
kernel
            [([Node]
dims [Node] -> Node -> Node
forall a. HasCallStack => [a] -> Node -> a
P.!! Node
x, Node -> Node
sf Node
x, Node -> Node
sg Node
x) | Node
x <- [Node]
keep]
            ((Node -> Node) -> [Node] -> Vector Node
offsets Node -> Node
sf [Node]
summed)
            ((Node -> Node) -> [Node] -> Vector Node
offsets Node -> Node
sg [Node]
summed)
            (Factor e -> Vector e
forall e. Factor e -> Vector e
values Factor e
f)
            (Factor e -> Vector e
forall e. Factor e -> Vector e
values Factor e
g)
        )
  where
    sf :: Node -> Node
sf = [Node] -> Factor e -> Node -> Node
forall e. [Node] -> Factor e -> Node -> Node
stride [Node]
dims Factor e
f
    sg :: Node -> Node
sg = [Node] -> Factor e -> Node -> Node
forall e. [Node] -> Factor e -> Node -> Node
stride [Node]
dims Factor e
g
    -- the offset into a factor's entries of each index of the given nodes
    offsets :: (Node -> Node) -> [Node] -> Vector Node
offsets Node -> Node
s [Node
x] = Node -> (Node -> Node) -> Vector Node
forall a. Unbox a => Node -> (Node -> a) -> Vector a
UV.generate ([Node]
dims [Node] -> Node -> Node
forall a. HasCallStack => [a] -> Node -> a
P.!! Node
x) (Node -> Node -> Node
forall a. Num a => a -> a -> a
* Node -> Node
s Node
x)
    offsets Node -> Node
s [Node]
xs = (Vector Node -> Node -> Vector Node)
-> Vector Node -> [Node] -> Vector Node
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
P.foldl (\Vector Node
acc Node
x -> Vector Node -> Node -> Node -> Vector Node
step Vector Node
acc ([Node]
dims [Node] -> Node -> Node
forall a. HasCallStack => [a] -> Node -> a
P.!! Node
x) (Node -> Node
s Node
x)) (Node -> Vector Node
forall a. Unbox a => a -> Vector a
UV.singleton Node
0) [Node]
xs
    step :: Vector Node -> Node -> Node -> Vector Node
step Vector Node
acc Node
d Node
s = Node -> (Node -> Node) -> Vector Node
forall a. Unbox a => Node -> (Node -> a) -> Vector a
UV.generate (Vector Node -> Node
forall a. Unbox a => Vector a -> Node
UV.length Vector Node
acc Node -> Node -> Node
forall a. Num a => a -> a -> a
* Node
d) \Node
j -> let (Node
q, Node
r) = Node
j Node -> Node -> (Node, Node)
forall a. Integral a => a -> a -> (a, a)
`P.quotRem` Node
d in Vector Node
acc Vector Node -> Node -> Node
forall a. Unbox a => Vector a -> Node -> a
`UV.unsafeIndex` Node
q Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
r Node -> Node -> Node
forall a. Num a => a -> a -> a
* Node
s

-- | One factor with the given nodes kept, in that order, and the others summed: its product with the
-- unit. A factor with an axis twice gives its diagonal, and the entries are the same along a kept
-- node that it does not have.
contract1 :: (Scalar e) => [Dim] -> Factor e -> [Node] -> [Node] -> Factor e
contract1 :: forall e.
Scalar e =>
[Node] -> Factor e -> [Node] -> [Node] -> Factor e
contract1 [Node]
dims Factor e
f [Node]
keep [Node]
summed
  | [Node] -> Bool
forall a. [a] -> Bool
forall (t :: Type -> Type) a. Foldable t => t a -> Bool
P.null [Node]
summed Bool -> Bool -> Bool
P.&& Factor e -> [Node]
forall e. Factor e -> [Node]
axes Factor e
f [Node] -> [Node] -> Bool
forall a. Eq a => a -> a -> Bool
== [Node]
keep = Factor e
f
  | Bool
P.otherwise = [Node] -> Factor e -> Factor e -> [Node] -> [Node] -> Factor e
forall e.
Scalar e =>
[Node] -> Factor e -> Factor e -> [Node] -> [Node] -> Factor e
contract [Node]
dims Factor e
f Factor e
forall e. Scalar e => Factor e
unit [Node]
keep [Node]
summed

-- | The factor without axes whose entry is 1.
unit :: (Scalar e) => Factor e
unit :: forall e. Scalar e => Factor e
unit = [Node] -> Vector e -> Factor e
forall e. Scalar e => [Node] -> Vector e -> Factor e
factor [] (e -> Vector e
forall a. Storable a => a -> Vector a
SV.singleton e
1)

-- | A contraction of two factors as a matrix product, when it is one: something is summed, every
-- summed node is on both factors, and the kept nodes of the first come before those of the second,
-- or after them, as the transposed product. A factor whose axes are in neither the order of its
-- matrix nor that of its transpose is rearranged first.
viaGemm :: (Scalar e) => Gemm e -> [Dim] -> Factor e -> Factor e -> [Node] -> [Node] -> P.Maybe (Factor e)
viaGemm :: forall e.
Scalar e =>
Gemm e
-> [Node]
-> Factor e
-> Factor e
-> [Node]
-> [Node]
-> Maybe (Factor e)
viaGemm Gemm e
mm [Node]
dims Factor e
f Factor e
g [Node]
keep [Node]
summed
  | Bool -> Bool
P.not (Factor e -> Factor e -> [Node] -> Bool
forall e. Factor e -> Factor e -> [Node] -> Bool
matrixProduct Factor e
f Factor e
g [Node]
summed) Bool -> Bool -> Bool
P.|| Node
rows Node -> Node -> Node
forall a. Num a => a -> a -> a
* Node
cols Node -> Node -> Node
forall a. Num a => a -> a -> a
* Node
inner Node -> Node -> Bool
forall a. Ord a => a -> a -> Bool
P.< Node
gemmWork = Maybe (Factor e)
forall a. Maybe a
P.Nothing
  | [Node]
keep [Node] -> [Node] -> Bool
forall a. Eq a => a -> a -> Bool
== [Node]
ms [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ [Node]
ns = Factor e -> Maybe (Factor e)
forall a. a -> Maybe a
P.Just ([Node] -> Vector e -> Factor e
forall e. Scalar e => [Node] -> Vector e -> Factor e
factor [Node]
keep (Gemm e
mm Bool
ta Bool
tb Node
rows Node
cols Node
inner Vector e
va Vector e
vb))
  | [Node]
keep [Node] -> [Node] -> Bool
forall a. Eq a => a -> a -> Bool
== [Node]
ns [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ [Node]
ms = Factor e -> Maybe (Factor e)
forall a. a -> Maybe a
P.Just ([Node] -> Vector e -> Factor e
forall e. Scalar e => [Node] -> Vector e -> Factor e
factor [Node]
keep (Gemm e
mm (Bool -> Bool
P.not Bool
tb) (Bool -> Bool
P.not Bool
ta) Node
cols Node
rows Node
inner Vector e
vb Vector e
va))
  | Bool
P.otherwise = Maybe (Factor e)
forall a. Maybe a
P.Nothing
  where
    ms :: [Node]
ms = [Node
x | Node
x <- Factor e -> [Node]
forall e. Factor e -> [Node]
axes Factor e
f, Bool -> Bool
P.not (Node -> Factor e -> Bool
forall e. Node -> Factor e -> Bool
onFactor Node
x Factor e
g)]
    ns :: [Node]
ns = [Node
x | Node
x <- Factor e -> [Node]
forall e. Factor e -> [Node]
axes Factor e
g, Bool -> Bool
P.not (Node -> Factor e -> Bool
forall e. Node -> Factor e -> Bool
onFactor Node
x Factor e
f)]
    ss :: [Node]
ss = (Node -> Bool) -> [Node] -> [Node]
forall a. (a -> Bool) -> [a] -> [a]
P.filter (Node -> [Node] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`P.elem` [Node]
summed) (Factor e -> [Node]
forall e. Factor e -> [Node]
axes Factor e
f)
    rows :: Node
rows = [Node] -> [Node] -> Node
extent [Node]
dims [Node]
ms
    cols :: Node
cols = [Node] -> [Node] -> Node
extent [Node]
dims [Node]
ns
    inner :: Node
inner = [Node] -> [Node] -> Node
extent [Node]
dims [Node]
ss
    (Bool
ta, Vector e
va) = Factor e -> [Node] -> [Node] -> (Bool, Vector e)
matrix Factor e
f [Node]
ms [Node]
ss
    (Bool
tb, Vector e
vb) = Factor e -> [Node] -> [Node] -> (Bool, Vector e)
matrix Factor e
g [Node]
ss [Node]
ns
    -- a factor as a matrix with the first nodes as rows: transposed, or rearranged when it is not
    -- already in that order
    matrix :: Factor e -> [Node] -> [Node] -> (Bool, Vector e)
matrix Factor e
h [Node]
rs [Node]
cs
      | Factor e -> [Node]
forall e. Factor e -> [Node]
axes Factor e
h [Node] -> [Node] -> Bool
forall a. Eq a => a -> a -> Bool
== [Node]
cs [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ [Node]
rs = (Bool
P.True, Factor e -> Vector e
forall e. Factor e -> Vector e
values Factor e
h)
      | Bool
P.otherwise = (Bool
P.False, Factor e -> Vector e
forall e. Factor e -> Vector e
values ([Node] -> Factor e -> [Node] -> [Node] -> Factor e
forall e.
Scalar e =>
[Node] -> Factor e -> [Node] -> [Node] -> Factor e
contract1 [Node]
dims Factor e
h ([Node]
rs [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ [Node]
cs) []))

-- | How far a step along a node moves in a factor's entries: the strides of its axes on that node,
-- 0 when it has none. Applied to a factor alone, it works out the strides of its axes once.
stride :: [Dim] -> Factor e -> Node -> Int
stride :: forall e. [Node] -> Factor e -> Node -> Node
stride [Node]
dims Factor e
f = \Node
node -> [Node] -> Node
forall a. Num a => [a] -> a
forall (t :: Type -> Type) a. (Foldable t, Num a) => t a -> a
P.sum [Node
s | (Node
a, Node
s) <- [(Node, Node)]
axisStrides, Node
a Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
== Node
node]
  where
    axisStrides :: [(Node, Node)]
axisStrides = [Node] -> [Node] -> [(Node, Node)]
forall a b. [a] -> [b] -> [(a, b)]
P.zip (Factor e -> [Node]
forall e. Factor e -> [Node]
axes Factor e
f) (Node -> [Node] -> [Node]
forall a. Node -> [a] -> [a]
P.drop Node
1 ((Node -> Node -> Node) -> Node -> [Node] -> [Node]
forall a b. (a -> b -> b) -> b -> [a] -> [b]
P.scanr Node -> Node -> Node
forall a. Num a => a -> a -> a
(*) Node
1 ((Node -> Node) -> [Node] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap ([Node]
dims [Node] -> Node -> Node
forall a. HasCallStack => [a] -> Node -> a
P.!!) (Factor e -> [Node]
forall e. Factor e -> [Node]
axes Factor e
f))))

-- | Whether no element of the list is there twice.
distinct :: [Node] -> P.Bool
distinct :: [Node] -> Bool
distinct = IntSet -> [Node] -> Bool
go IntSet
IS.empty
  where
    go :: IntSet -> [Node] -> Bool
go IntSet
_ [] = Bool
P.True
    go IntSet
seen (Node
x : [Node]
xs) = Bool -> Bool
P.not (Node
x Node -> IntSet -> Bool
`IS.member` IntSet
seen) Bool -> Bool -> Bool
P.&& IntSet -> [Node] -> Bool
go (Node -> IntSet -> IntSet
IS.insert Node
x IntSet
seen) [Node]
xs

-- | Whether the node is on one of the factor's axes.
onFactor :: Node -> Factor e -> P.Bool
onFactor :: forall e. Node -> Factor e -> Bool
onFactor Node
x Factor e
f = Node
x Node -> [Node] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: Type -> Type) a.
(Foldable t, Eq a) =>
a -> t a -> Bool
`P.elem` Factor e -> [Node]
forall e. Factor e -> [Node]
axes Factor e
f

-- | Whether a product of two factors that sums the given nodes is a matrix product: something is
-- summed, and the nodes on both factors are the summed ones.
matrixProduct :: Factor e -> Factor e -> [Node] -> P.Bool
matrixProduct :: forall e. Factor e -> Factor e -> [Node] -> Bool
matrixProduct Factor e
f Factor e
g [Node]
summed = Bool -> Bool
P.not ([Node] -> Bool
forall a. [a] -> Bool
forall (t :: Type -> Type) a. Foldable t => t a -> Bool
P.null [Node]
summed) Bool -> Bool -> Bool
P.&& [Node] -> Set Node
forall a. Ord a => [a] -> Set a
Set.fromList [Node
x | Node
x <- Factor e -> [Node]
forall e. Factor e -> [Node]
axes Factor e
f, Node -> Factor e -> Bool
forall e. Node -> Factor e -> Bool
onFactor Node
x Factor e
g] Set Node -> Set Node -> Bool
forall a. Eq a => a -> a -> Bool
== [Node] -> Set Node
forall a. Ord a => [a] -> Set a
Set.fromList [Node]
summed

-- | The number of indices of the given nodes together.
extent :: [Dim] -> [Node] -> Int
extent :: [Node] -> [Node] -> Node
extent [Node]
dims = [Node] -> Node
forall a. Num a => [a] -> a
forall (t :: Type -> Type) a. (Foldable t, Num a) => t a -> a
P.product ([Node] -> Node) -> ([Node] -> [Node]) -> [Node] -> Node
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. (Node -> Node) -> [Node] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap ([Node]
dims [Node] -> Node -> Node
forall a. HasCallStack => [a] -> Node -> a
P.!!)

-- | The number of multiplications from which a matrix product goes to 'gemm'.
gemmWork :: Int
gemmWork :: Node
gemmWork = Node
4096

-- | The product of two matrices given row by row, for products that are not part of a network:
-- whether the first and the second are stored transposed, the numbers of rows, columns and summed
-- indices, and the two matrices; the result row by row. It is the contraction of two factors on the
-- nodes of the rows, the columns and the summed index, so large products go to 'gemm'.
matMul :: (Scalar e) => P.Bool -> P.Bool -> Int -> Int -> Int -> SV.Vector e -> SV.Vector e -> SV.Vector e
matMul :: forall e.
Scalar e =>
Bool
-> Bool -> Node -> Node -> Node -> Vector e -> Vector e -> Vector e
matMul Bool
ta Bool
tb Node
rows Node
cols Node
inner Vector e
a Vector e
b =
  Factor e -> Vector e
forall e. Factor e -> Vector e
values
    ( [Node] -> Factor e -> Factor e -> [Node] -> [Node] -> Factor e
forall e.
Scalar e =>
[Node] -> Factor e -> Factor e -> [Node] -> [Node] -> Factor e
contract
        [Node
rows, Node
cols, Node
inner]
        ([Node] -> Vector e -> Factor e
forall e. Scalar e => [Node] -> Vector e -> Factor e
factor (if Bool
ta then [Node
2, Node
0] else [Node
0, Node
2]) Vector e
a)
        ([Node] -> Vector e -> Factor e
forall e. Scalar e => [Node] -> Vector e -> Factor e
factor (if Bool
tb then [Node
1, Node
2] else [Node
2, Node
1]) Vector e
b)
        [Node
0, Node
1]
        [Node
2]
    )

-- | The network with every factor's axes distinct, every node that is on no wire summed out, and
-- the nodes renumbered.
normalize :: (Scalar e) => Net e -> Net e
normalize :: forall e. Scalar e => Net e -> Net e
normalize (Net [Node]
dims [Node]
is [Node]
os [Factor e]
fs) =
  [Node] -> [Node] -> [Node] -> [Factor e] -> Net e
forall e. [Node] -> [Node] -> [Node] -> [Factor e] -> Net e
Net
    ((Node -> Node) -> [Node] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap ([Node]
dims [Node] -> Node -> Node
forall a. HasCallStack => [a] -> Node -> a
P.!!) [Node]
kept)
    ((Node -> Node) -> [Node] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap Node -> Node
renumber [Node]
is)
    ((Node -> Node) -> [Node] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap Node -> Node
renumber [Node]
os)
    ((Factor e -> Factor e) -> [Factor e] -> [Factor e]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap Factor e -> Factor e
relabel ([Factor e] -> [Factor e]
forall {e}. Scalar e => [Factor e] -> [Factor e]
scalars ([Factor e]
loops [Factor e] -> [Factor e] -> [Factor e]
forall a. [a] -> [a] -> [a]
P.++ (Plan -> Factor e) -> [Plan] -> [Factor e]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap ((Factor e, IntMap Node) -> Factor e
forall a b. (a, b) -> a
P.fst ((Factor e, IntMap Node) -> Factor e)
-> (Plan -> (Factor e, IntMap Node)) -> Plan -> Factor e
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. Plan -> (Factor e, IntMap Node)
run) [Plan]
plans)))
  where
    summed :: [Node]
summed = [Node
n | Node
n <- [Node
0 .. [Node] -> Node
forall a. [a] -> Node
forall (t :: Type -> Type) a. Foldable t => t a -> Node
P.length [Node]
dims Node -> Node -> Node
forall a. Num a => a -> a -> a
- Node
1], Node
n Node -> Set Node -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` Set Node
keptSet]
    -- the number of factors on each node
    factorsOn :: Node -> Node
factorsOn Node
n = Node -> Node -> IntMap Node -> Node
forall a. a -> Node -> IntMap a -> a
IM.findWithDefault Node
0 Node
n IntMap Node
counts
    distinctAxes :: [(Factor e, [Node])]
distinctAxes = [(Factor e
f, [Node] -> [Node]
forall a. Ord a => [a] -> [a]
nubOrd (Factor e -> [Node]
forall e. Factor e -> [Node]
axes Factor e
f)) | Factor e
f <- [Factor e]
fs]
    counts :: IntMap Node
counts = (Node -> Node -> Node) -> [(Node, Node)] -> IntMap Node
forall a. (a -> a -> a) -> [(Node, a)] -> IntMap a
IM.fromListWith Node -> Node -> Node
forall a. Num a => a -> a -> a
(+) [(Node
m, Node
1 :: Int) | (Factor e
_, [Node]
as) <- [(Factor e, [Node])]
distinctAxes, Node
m <- [Node]
as]
    -- each factor with its axes distinct, and the summed nodes that no other factor has summed out
    reduced :: [Factor e]
reduced =
      [ let own :: [Node]
own = [Node
m | Node
m <- [Node]
as, Node
m Node -> Set Node -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.notMember` Set Node
keptSet, Node -> Node
factorsOn Node
m Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
== Node
1] in [Node] -> Factor e -> [Node] -> [Node] -> Factor e
forall e.
Scalar e =>
[Node] -> Factor e -> [Node] -> [Node] -> Factor e
contract1 [Node]
dims Factor e
f ([Node]
as [Node] -> [Node] -> [Node]
forall a. Eq a => [a] -> [a] -> [a]
List.\\ [Node]
own) [Node]
own
      | (Factor e
f, [Node]
as) <- [(Factor e, [Node])]
distinctAxes
      ]
    -- the summed nodes on no factor, as of closed loops: a scalar of their dimensions
    loops :: [Factor e]
loops = case [Node
n | Node
n <- [Node]
summed, Node -> Node
factorsOn Node
n Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
== Node
0] of
      [] -> []
      [Node]
ns -> [[Node] -> Vector e -> Factor e
forall e. Scalar e => [Node] -> Vector e -> Factor e
factor [] (e -> Vector e
forall a. Storable a => a -> Vector a
SV.singleton (Node -> e
forall a b. (Integral a, Num b) => a -> b
P.fromIntegral ([Node] -> [Node] -> Node
extent [Node]
dims [Node]
ns)))]
    shared :: [Node]
shared = [Node
n | Node
n <- [Node]
summed, Node -> Node
factorsOn Node
n Node -> Node -> Bool
forall a. Ord a => a -> a -> Bool
P.> Node
1]
    sharedSet :: Set Node
sharedSet = [Node] -> Set Node
forall a. Ord a => [a] -> Set a
Set.fromList [Node]
shared
    -- the factors contracted in the order of 'Plan.planContraction', until no two share a summed node
    plans :: [Plan]
plans =
      (Set Node -> Double)
-> (Node -> Bool) -> Set Node -> [Set Node] -> [Plan]
forall n.
Ord n =>
(Set n -> Double) -> (n -> Bool) -> Set n -> [Set n] -> [Plan]
Plan.planContraction
        (Node -> Double
forall a b. (Integral a, Num b) => a -> b
P.fromIntegral (Node -> Double) -> (Set Node -> Node) -> Set Node -> Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. [Node] -> [Node] -> Node
extent [Node]
dims ([Node] -> Node) -> (Set Node -> [Node]) -> Set Node -> Node
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. Set Node -> [Node]
forall a. Set a -> [a]
Set.toList)
        (Node -> Set Node -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set Node
sharedSet)
        Set Node
keptSet
        ((Factor e -> Set Node) -> [Factor e] -> [Set Node]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap ([Node] -> Set Node
forall a. Ord a => [a] -> Set a
Set.fromList ([Node] -> Set Node)
-> (Factor e -> [Node]) -> Factor e -> Set Node
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. Factor e -> [Node]
forall e. Factor e -> [Node]
axes) [Factor e]
reduced)
    reducedAt :: IntMap (Factor e)
reducedAt = [(Node, Factor e)] -> IntMap (Factor e)
forall a. [(Node, a)] -> IntMap a
IM.fromList ([Node] -> [Factor e] -> [(Node, Factor e)]
forall a b. [a] -> [b] -> [(a, b)]
P.zip [Node
0 ..] [Factor e]
reduced)
    -- the factor a plan makes, with how many of the factors in it touch each node
    run :: Plan -> (Factor e, IntMap Node)
run (Plan.Factor Node
k) = let f :: Factor e
f = IntMap (Factor e)
reducedAt IntMap (Factor e) -> Node -> Factor e
forall a. IntMap a -> Node -> a
IM.! Node
k in (Factor e
f, [(Node, Node)] -> IntMap Node
forall a. [(Node, a)] -> IntMap a
IM.fromList [(Node
m, Node
1 :: Int) | Node
m <- Factor e -> [Node]
forall e. Factor e -> [Node]
axes Factor e
f])
    run (Plan.Contract Plan
a Plan
b) =
      let (Factor e
f, IntMap Node
cf) = Plan -> (Factor e, IntMap Node)
run Plan
a
          (Factor e
g, IntMap Node
cg) = Plan -> (Factor e, IntMap Node)
run Plan
b
          cs :: IntMap Node
cs = (Node -> Node -> Node) -> IntMap Node -> IntMap Node -> IntMap Node
forall a. (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
IM.unionWith Node -> Node -> Node
forall a. Num a => a -> a -> a
(+) IntMap Node
cf IntMap Node
cg
          both :: [Node]
both = [Node] -> [Node]
forall a. Ord a => [a] -> [a]
nubOrd (Factor e -> [Node]
forall e. Factor e -> [Node]
axes Factor e
f [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ Factor e -> [Node]
forall e. Factor e -> [Node]
axes Factor e
g)
          -- the shared nodes that no factor outside the two is on, summed out by this product
          ds :: [Node]
ds = [Node
m | Node
m <- [Node]
both, Node
m Node -> Set Node -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set Node
sharedSet, IntMap Node
cs IntMap Node -> Node -> Node
forall a. IntMap a -> Node -> a
IM.! Node
m Node -> Node -> Bool
forall a. Eq a => a -> a -> Bool
== Node -> Node
factorsOn Node
m]
          left :: [Node]
left = [Node]
both [Node] -> [Node] -> [Node]
forall a. Eq a => [a] -> [a] -> [a]
List.\\ [Node]
ds
          -- a product that is not a matrix product puts the nodes that a later product sums last,
          -- where a matrix product wants them
          keep :: [Node]
keep =
            if Factor e -> Factor e -> [Node] -> Bool
forall e. Factor e -> Factor e -> [Node] -> Bool
matrixProduct Factor e
f Factor e
g [Node]
ds then [Node]
left else let ([Node]
later, [Node]
now) = (Node -> Bool) -> [Node] -> ([Node], [Node])
forall a. (a -> Bool) -> [a] -> ([a], [a])
List.partition (Node -> Set Node -> Bool
forall a. Ord a => a -> Set a -> Bool
`Set.member` Set Node
sharedSet) [Node]
left in [Node]
now [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ [Node]
later
      in ([Node] -> Factor e -> Factor e -> [Node] -> [Node] -> Factor e
forall e.
Scalar e =>
[Node] -> Factor e -> Factor e -> [Node] -> [Node] -> Factor e
contract [Node]
dims Factor e
f Factor e
g [Node]
keep [Node]
ds, IntMap Node
cs)
    -- the factors without axes multiplied into one
    scalars :: [Factor e] -> [Factor e]
scalars [Factor e]
gs = case (Factor e -> Bool) -> [Factor e] -> ([Factor e], [Factor e])
forall a. (a -> Bool) -> [a] -> ([a], [a])
List.partition ([Node] -> Bool
forall a. [a] -> Bool
forall (t :: Type -> Type) a. Foldable t => t a -> Bool
P.null ([Node] -> Bool) -> (Factor e -> [Node]) -> Factor e -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. Factor e -> [Node]
forall e. Factor e -> [Node]
axes) [Factor e]
gs of
      (Factor e
c : Factor e
d : [Factor e]
more, [Factor e]
rest) -> [Node] -> Vector e -> Factor e
forall e. Scalar e => [Node] -> Vector e -> Factor e
factor [] (e -> Vector e
forall a. Storable a => a -> Vector a
SV.singleton ([e] -> e
forall a. Num a => [a] -> a
forall (t :: Type -> Type) a. (Foldable t, Num a) => t a -> a
P.product [Vector e -> e
forall a. Storable a => Vector a -> a
SV.head (Factor e -> Vector e
forall e. Factor e -> Vector e
values Factor e
x) | Factor e
x <- Factor e
c Factor e -> [Factor e] -> [Factor e]
forall a. a -> [a] -> [a]
: Factor e
d Factor e -> [Factor e] -> [Factor e]
forall a. a -> [a] -> [a]
: [Factor e]
more])) Factor e -> [Factor e] -> [Factor e]
forall a. a -> [a] -> [a]
: [Factor e]
rest
      ([Factor e], [Factor e])
_ -> [Factor e]
gs
    kept :: [Node]
kept = [Node] -> [Node]
forall a. Ord a => [a] -> [a]
nubOrd ([Node]
is [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ [Node]
os)
    keptSet :: Set Node
keptSet = [Node] -> Set Node
forall a. Ord a => [a] -> Set a
Set.fromList [Node]
kept
    renumber :: Node -> Node
renumber = [Node] -> Node -> Node
renumbering [Node]
kept
    relabel :: Factor e -> Factor e
relabel Factor e
f = Factor e
f{axes = P.fmap renumber (axes f)}

-- | The new number of each of the given nodes: its place in the list.
renumbering :: [Node] -> Node -> Node
renumbering :: [Node] -> Node -> Node
renumbering [Node]
ns = ([(Node, Node)] -> IntMap Node
forall a. [(Node, a)] -> IntMap a
IM.fromList ([Node] -> [Node] -> [(Node, Node)]
forall a b. [a] -> [b] -> [(a, b)]
P.zip [Node]
ns [Node
0 ..]) IntMap Node -> Node -> Node
forall a. IntMap a -> Node -> a
IM.!)

-- | Two networks side by side.
besideNet :: Net e -> Net e -> Net e
besideNet :: forall e. Net e -> Net e -> Net e
besideNet (Net [Node]
d1 [Node]
i1 [Node]
o1 [Factor e]
f1) (Net [Node]
d2 [Node]
i2 [Node]
o2 [Factor e]
f2) =
  [Node] -> [Node] -> [Node] -> [Factor e] -> Net e
forall e. [Node] -> [Node] -> [Node] -> [Factor e] -> Net e
Net ([Node]
d1 [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ [Node]
d2) ([Node]
i1 [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ (Node -> Node) -> [Node] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap (Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
k) [Node]
i2) ([Node]
o1 [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ (Node -> Node) -> [Node] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap (Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
k) [Node]
o2) ([Factor e]
f1 [Factor e] -> [Factor e] -> [Factor e]
forall a. [a] -> [a] -> [a]
P.++ (Factor e -> Factor e) -> [Factor e] -> [Factor e]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap Factor e -> Factor e
shift [Factor e]
f2)
  where
    k :: Node
k = [Node] -> Node
forall a. [a] -> Node
forall (t :: Type -> Type) a. Foldable t => t a -> Node
P.length [Node]
d1
    shift :: Factor e -> Factor e
shift Factor e
f = Factor e
f{axes = P.fmap (+ k) (axes f)}

-- | The second network after the first: their wirings glued along the wires in between, and the
-- nodes that are then on no wire summed out.
composeNet :: (Scalar e) => Net e -> Net e -> Net e
composeNet :: forall e. Scalar e => Net e -> Net e -> Net e
composeNet Net e
f Net e
g
  | P.Just [Node]
through <- Net e -> Maybe [Node]
forall e. Net e -> Maybe [Node]
permutation Net e
g = Net e
f{outs = P.fmap (outs f P.!!) through}
  | P.Just [Node]
through <- Net e -> Maybe [Node]
forall e. Net e -> Maybe [Node]
permutation (Net e -> Net e
forall e. Net e -> Net e
transposeNet Net e
f) = Net e
g{ins = P.fmap (ins g P.!!) through}
  | Bool
P.otherwise =
      Net e -> Net e
forall e. Scalar e => Net e -> Net e
normalize
        ( [Node] -> [Node] -> [Node] -> [Factor e] -> Net e
forall e. [Node] -> [Node] -> [Node] -> [Factor e] -> Net e
Net
            [Net e -> [Node]
forall e. Net e -> [Node]
nodeDims Net e
both [Node] -> Node -> Node
forall a. HasCallStack => [a] -> Node -> a
P.!! Node
r | Node
r <- [Node]
live]
            ((Node -> Node) -> [Node] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap Node -> Node
node (Net e -> [Node]
forall e. Net e -> [Node]
ins Net e
f))
            ((Node -> Node) -> [Node] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap (Node -> Node
node (Node -> Node) -> (Node -> Node) -> Node -> Node
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. (Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
k)) (Net e -> [Node]
forall e. Net e -> [Node]
outs Net e
g))
            [Factor e
h{axes = P.fmap node (axes h)} | Factor e
h <- Net e -> [Factor e]
forall e. Net e -> [Factor e]
factors Net e
both]
        )
  where
    both :: Net e
both = Net e -> Net e -> Net e
forall e. Net e -> Net e -> Net e
besideNet Net e
f Net e
g
    k :: Node
k = [Node] -> Node
forall a. [a] -> Node
forall (t :: Type -> Type) a. Foldable t => t a -> Node
P.length (Net e -> [Node]
forall e. Net e -> [Node]
nodeDims Net e
f)
    find :: Node -> Node
find = [(Node, Node)] -> Node -> Node
forall a. Ord a => [(a, a)] -> a -> a
unionFind ([Node] -> [Node] -> [(Node, Node)]
forall a b. [a] -> [b] -> [(a, b)]
P.zip (Net e -> [Node]
forall e. Net e -> [Node]
outs Net e
f) ((Node -> Node) -> [Node] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap (Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
k) (Net e -> [Node]
forall e. Net e -> [Node]
ins Net e
g)))
    -- the nodes that are left after gluing, numbered from 0
    live :: [Node]
live = [Node] -> [Node]
forall a. Ord a => [a] -> [a]
nubOrd [Node -> Node
find Node
n | Node
n <- [Node
0 .. [Node] -> Node
forall a. [a] -> Node
forall (t :: Type -> Type) a. Foldable t => t a -> Node
P.length (Net e -> [Node]
forall e. Net e -> [Node]
nodeDims Net e
both) Node -> Node -> Node
forall a. Num a => a -> a -> a
- Node
1]]
    node :: Node -> Node
node = [Node] -> Node -> Node
renumbering [Node]
live (Node -> Node) -> (Node -> Node) -> Node -> Node
forall b c a. (b -> c) -> (a -> b) -> a -> c
P.. Node -> Node
find

-- | For a wiring without factors whose every node has one input wire and at least one output wire,
-- such as a permutation or copying: the input wire that each output wire continues.
permutation :: Net e -> P.Maybe [Int]
permutation :: forall e. Net e -> Maybe [Node]
permutation (Net [Node]
_ [Node]
is [Node]
os [Factor e]
fs)
  | Bool -> Bool
P.not ([Factor e] -> Bool
forall a. [a] -> Bool
forall (t :: Type -> Type) a. Foldable t => t a -> Bool
P.null [Factor e]
fs) Bool -> Bool -> Bool
P.|| Bool -> Bool
P.not ([Node] -> Bool
distinct [Node]
is) = Maybe [Node]
forall a. Maybe a
P.Nothing
  | [Node]
is [Node] -> [Node] -> Bool
forall a. Eq a => a -> a -> Bool
== [Node]
os = [Node] -> Maybe [Node]
forall a. a -> Maybe a
P.Just [Node
0 .. [Node] -> Node
forall a. [a] -> Node
forall (t :: Type -> Type) a. Foldable t => t a -> Node
P.length [Node]
is Node -> Node -> Node
forall a. Num a => a -> a -> a
- Node
1]
  | [Node] -> Set Node
forall a. Ord a => [a] -> Set a
Set.fromList [Node]
is Set Node -> Set Node -> Bool
forall a. Eq a => a -> a -> Bool
== [Node] -> Set Node
forall a. Ord a => [a] -> Set a
Set.fromList [Node]
os = [Node] -> Maybe [Node]
forall a. a -> Maybe a
P.Just ((Node -> Node) -> [Node] -> [Node]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap ([Node] -> Node -> Node
renumbering [Node]
is) [Node]
os)
  | Bool
P.otherwise = Maybe [Node]
forall a. Maybe a
P.Nothing

-- | The network the other way round.
transposeNet :: Net e -> Net e
transposeNet :: forall e. Net e -> Net e
transposeNet (Net [Node]
d [Node]
i [Node]
o [Factor e]
fs) = [Node] -> [Node] -> [Node] -> [Factor e] -> Net e
forall e. [Node] -> [Node] -> [Node] -> [Factor e] -> Net e
Net [Node]
d [Node]
o [Node]
i [Factor e]
fs

withAppend
  :: forall as bs r. (KnownListOf KnownNat as, KnownListOf KnownNat bs) => ((KnownListOf KnownNat (as ++ bs)) => r) -> r
withAppend :: forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend = forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs) =>
(KnownListOf c (as ++ bs) => r) -> r
withKnownAppend @KnownNat @as @bs

-- | The arrow the other way round: its wiring with inputs and outputs swapped.
transpose :: TensorNetwork (a :: TNET e) b -> TensorNetwork b a
transpose :: forall e (a :: TNET e) (b :: TNET e).
TensorNetwork a b -> TensorNetwork b a
transpose (TensorNetwork Net e
n) = Net e -> TensorNetwork (TN bs) (TN as)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Net e -> TensorNetwork (TN as) (TN bs)
TensorNetwork (Net e -> Net e
forall e. Net e -> Net e
transposeNet Net e
n)

instance (Scalar e) => Profunctor (TensorNetwork :: CAT (TNET e)) where
  dimap :: forall (c :: TNET e) (a :: TNET e) (b :: TNET e) (d :: TNET e).
(c ~> a) -> (b ~> d) -> TensorNetwork a b -> TensorNetwork c d
dimap = (c ~> a) -> (b ~> d) -> TensorNetwork a b -> TensorNetwork c d
TensorNetwork c a
-> TensorNetwork b d -> TensorNetwork a b -> TensorNetwork c d
forall {k} (p :: CAT k) (c :: k) (a :: k) (b :: k) (d :: k).
Promonad p =>
p c a -> p b d -> p a b -> p c d
dimapDefault
  (Ob a, Ob b) => r
r \\ :: forall (a :: TNET e) (b :: TNET e) r.
((Ob a, Ob b) => r) -> TensorNetwork a b -> r
\\ TensorNetwork{} = r
(Ob a, Ob b) => r
r

instance (Scalar e) => Promonad (TensorNetwork :: CAT (TNET e)) where
  id :: forall (a :: TNET e). Ob a => TensorNetwork a a
id = TensorNetwork a a
TensorNetwork (TN (UN TN a)) (TN (UN TN a))
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
TensorNetwork (TN as) (TN bs)
straight
  TensorNetwork Net e
g . :: forall (b :: TNET e) (c :: TNET e) (a :: TNET e).
TensorNetwork b c -> TensorNetwork a b -> TensorNetwork a c
. TensorNetwork Net e
f = Net e -> TensorNetwork (TN as) (TN bs)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Net e -> TensorNetwork (TN as) (TN bs)
TensorNetwork (Net e -> Net e -> Net e
forall e. Scalar e => Net e -> Net e -> Net e
composeNet Net e
f Net e
g)

-- | Tensor networks with entries @e@, between lists of dimensions.
instance (Scalar e) => CategoryOf (TNET e) where
  type (~>) = TensorNetwork
  type Ob a = (Is TN a, KnownListOf KnownNat (UN TN a))

-- | The conjugate transpose: the factors with their entries and conjugates swapped. 'dual' is the
-- transpose without conjugating, since the compact-closed structure is bilinear.
instance (Scalar e) => DaggerProfunctor (TensorNetwork :: CAT (TNET e)) where
  dagger :: forall (a :: TNET e) (b :: TNET e).
TensorNetwork a b -> TensorNetwork b a
dagger (TensorNetwork Net e
n) = Net e -> TensorNetwork (TN bs) (TN as)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Net e -> TensorNetwork (TN as) (TN bs)
TensorNetwork (Net e -> Net e
forall e. Net e -> Net e
transposeNet Net e
n){factors = [h{values = conjugates h, conjugates = values h} | h <- factors n]}

-- | A wire of dimension 0.
instance (Scalar e) => HasInitialObject (TNET e) where
  type InitialObject = TN '[0]
  initiate :: forall (a :: TNET e). Ob a => InitialObject ~> a
initiate = InitialObject ~> a
TensorNetwork (TN '[0]) (TN (UN TN a))
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
TensorNetwork (TN as) (TN bs)
zero

-- | A wire of dimension 0.
instance (Scalar e) => HasTerminalObject (TNET e) where
  type TerminalObject = TN '[0]
  terminate :: forall (a :: TNET e). Ob a => a ~> TerminalObject
terminate = a ~> TerminalObject
TensorNetwork (TN (UN TN a)) (TN '[0])
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
TensorNetwork (TN as) (TN bs)
zero

-- | The direct sum, as a single wire: the transpose of the product.
instance (Scalar e) => HasBinaryCoproducts (TNET e) where
  type a || b = TN '[Size (UN TN a) + Size (UN TN b)]
  withObCoprod :: forall (a :: TNET e) (b :: TNET e) r.
(Ob a, Ob b) =>
(Ob (a || b) => r) -> r
withObCoprod @(TN as) @(TN bs) Ob (a || b) => r
r = forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownNat (Size as + Size bs) => r) -> r
withSum @as @bs r
KnownNat (Size (UN TN a) + Size (UN TN b)) => r
Ob (a || b) => r
r
  lft :: forall (a :: TNET e) (b :: TNET e). (Ob a, Ob b) => a ~> (a || b)
lft @(TN as) @(TN bs) = forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownNat (Size as + Size bs) => r) -> r
withSum @as @bs ((Node -> Node)
-> TensorNetwork
     (TN (UN TN a)) (TN '[Size (UN TN a) + Size (UN TN b)])
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Node -> Node) -> TensorNetwork (TN as) (TN bs)
reindex Node -> Node
forall a. a -> a
P.id)
  rgt :: forall (a :: TNET e) (b :: TNET e). (Ob a, Ob b) => b ~> (a || b)
rgt @(TN as) @(TN bs) = forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownNat (Size as + Size bs) => r) -> r
withSum @as @bs ((Node -> Node)
-> TensorNetwork
     (TN (UN TN b)) (TN '[Size (UN TN a) + Size (UN TN b)])
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Node -> Node) -> TensorNetwork (TN as) (TN bs)
reindex (forall (ns :: [Nat]). KnownListOf KnownNat ns => Node
size @as Node -> Node -> Node
forall a. Num a => a -> a -> a
+))
  x ~> a
f ||| :: forall (x :: TNET e) (a :: TNET e) (y :: TNET e).
(x ~> a) -> (y ~> a) -> (x || y) ~> a
||| y ~> a
g = TensorNetwork a (TN '[Size (UN TN x) + Size (UN TN y)])
-> TensorNetwork (TN '[Size (UN TN x) + Size (UN TN y)]) a
forall e (a :: TNET e) (b :: TNET e).
TensorNetwork a b -> TensorNetwork b a
transpose (TensorNetwork x a -> TensorNetwork a x
forall e (a :: TNET e) (b :: TNET e).
TensorNetwork a b -> TensorNetwork b a
transpose x ~> a
TensorNetwork x a
f (a ~> x) -> (a ~> y) -> a ~> (x && y)
forall k (a :: k) (x :: k) (y :: k).
HasBinaryProducts k =>
(a ~> x) -> (a ~> y) -> a ~> (x && y)
forall (a :: TNET e) (x :: TNET e) (y :: TNET e).
(a ~> x) -> (a ~> y) -> a ~> (x && y)
&&& TensorNetwork y a -> TensorNetwork a y
forall e (a :: TNET e) (b :: TNET e).
TensorNetwork a b -> TensorNetwork b a
transpose y ~> a
TensorNetwork y a
g)

-- | The direct sum, as a single wire: the entries of the first arrow above those of the second.
instance (Scalar e) => HasBinaryProducts (TNET e) where
  type a && b = TN '[Size (UN TN a) + Size (UN TN b)]
  withObProd :: forall (a :: TNET e) (b :: TNET e) r.
(Ob a, Ob b) =>
(Ob (a && b) => r) -> r
withObProd @(TN as) @(TN bs) Ob (a && b) => r
r = forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownNat (Size as + Size bs) => r) -> r
withSum @as @bs r
KnownNat (Size (UN TN a) + Size (UN TN b)) => r
Ob (a && b) => r
r
  fst :: forall (a :: TNET e) (b :: TNET e). (Ob a, Ob b) => (a && b) ~> a
fst @a @b = TensorNetwork
  (TN (UN TN a)) (TN '[Size (UN TN a) + Size (UN TN b)])
-> TensorNetwork
     (TN '[Size (UN TN a) + Size (UN TN b)]) (TN (UN TN a))
forall e (a :: TNET e) (b :: TNET e).
TensorNetwork a b -> TensorNetwork b a
transpose (forall k (a :: k) (b :: k).
(HasBinaryCoproducts k, Ob a, Ob b) =>
a ~> (a || b)
lft @_ @a @b)
  snd :: forall (a :: TNET e) (b :: TNET e). (Ob a, Ob b) => (a && b) ~> b
snd @a @b = TensorNetwork
  (TN (UN TN b)) (TN '[Size (UN TN a) + Size (UN TN b)])
-> TensorNetwork
     (TN '[Size (UN TN a) + Size (UN TN b)]) (TN (UN TN b))
forall e (a :: TNET e) (b :: TNET e).
TensorNetwork a b -> TensorNetwork b a
transpose (forall k (a :: k) (b :: k).
(HasBinaryCoproducts k, Ob a, Ob b) =>
b ~> (a || b)
rgt @_ @a @b)
  f :: a ~> x
f@(TensorNetwork @_ @as Net e
_) &&& :: forall (a :: TNET e) (x :: TNET e) (y :: TNET e).
(a ~> x) -> (a ~> y) -> a ~> (x && y)
&&& g :: a ~> y
g@(TensorNetwork @_ @bs Net e
_) = forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownNat (Size as + Size bs) => r) -> r
withSum @as @bs (Vector e -> TensorNetwork (TN as) (TN '[Size bs + Size bs])
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Vector e -> TensorNetwork (TN as) (TN bs)
fromVector (TensorNetwork (TN as) (TN bs) -> Vector e
forall {e} (a :: TNET e) (b :: TNET e).
Scalar e =>
TensorNetwork a b -> Vector e
toVector a ~> x
TensorNetwork (TN as) (TN bs)
f Vector e -> Vector e -> Vector e
forall a. Storable a => Vector a -> Vector a -> Vector a
SV.++ TensorNetwork (TN as) (TN bs) -> Vector e
forall {e} (a :: TNET e) (b :: TNET e).
Scalar e =>
TensorNetwork a b -> Vector e
toVector a ~> y
TensorNetwork (TN as) (TN bs)
g))

instance (Scalar e) => HasBiproducts (TNET e)

instance (Scalar e) => MonoidalProfunctor (TensorNetwork :: CAT (TNET e)) where
  one :: TensorNetwork Unit Unit
one = TensorNetwork Unit Unit
TensorNetwork (TN '[]) (TN '[])
forall {k} (p :: CAT k) (a :: k). (Promonad p, Ob a) => p a a
forall (a :: TNET e). Ob a => TensorNetwork a a
id
  TensorNetwork @as @bs Net e
f ** :: forall (x1 :: TNET e) (x2 :: TNET e) (y1 :: TNET e) (y2 :: TNET e).
TensorNetwork x1 x2
-> TensorNetwork y1 y2 -> TensorNetwork (x1 ** y1) (x2 ** y2)
** TensorNetwork @cs @ds Net e
g = forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @as @cs (forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @bs @ds (Net e -> TensorNetwork (TN (as ++ as)) (TN (bs ++ bs))
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Net e -> TensorNetwork (TN as) (TN bs)
TensorNetwork (Net e -> Net e -> Net e
forall e. Net e -> Net e -> Net e
besideNet Net e
f Net e
g)))

-- | The wires side by side as the tensor: on matrices, the Kronecker product.
instance (Scalar e) => Monoidal (TNET e) where
  type Unit = TN '[]
  type a ** b = TN (UN TN a ++ UN TN b)
  withOb2 :: forall (a :: TNET e) (b :: TNET e) r.
(Ob a, Ob b) =>
(Ob (a ** b) => r) -> r
withOb2 @(TN as) @(TN bs) Ob (a ** b) => r
r = forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @as @bs r
Ob (a ** b) => r
KnownListOf KnownNat (UN TN a ++ UN TN b) => r
r
  leftUnitor :: forall (a :: TNET e). Ob a => (Unit ** a) ~> a
leftUnitor = (Unit ** a) ~> a
TensorNetwork (TN (UN TN a)) (TN (UN TN a))
forall {k} (p :: CAT k) (a :: k). (Promonad p, Ob a) => p a a
forall (a :: TNET e). Ob a => TensorNetwork a a
id
  leftUnitorInv :: forall (a :: TNET e). Ob a => a ~> (Unit ** a)
leftUnitorInv = a ~> (Unit ** a)
TensorNetwork (TN (UN TN a)) (TN (UN TN a))
forall {k} (p :: CAT k) (a :: k). (Promonad p, Ob a) => p a a
forall (a :: TNET e). Ob a => TensorNetwork a a
id
  rightUnitor :: forall (a :: TNET e). Ob a => (a ** Unit) ~> a
rightUnitor @(TN as) = forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @as @'[] TensorNetwork (TN (UN TN a ++ '[])) (TN (UN TN a))
KnownListOf KnownNat (UN TN a ++ '[]) =>
TensorNetwork (TN (UN TN a ++ '[])) (TN (UN TN a))
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
TensorNetwork (TN as) (TN bs)
straight
  rightUnitorInv :: forall (a :: TNET e). Ob a => a ~> (a ** Unit)
rightUnitorInv @(TN as) = forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @as @'[] TensorNetwork (TN (UN TN a)) (TN (UN TN a ++ '[]))
KnownListOf KnownNat (UN TN a ++ '[]) =>
TensorNetwork (TN (UN TN a)) (TN (UN TN a ++ '[]))
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
TensorNetwork (TN as) (TN bs)
straight
  associator :: forall (a :: TNET e) (b :: TNET e) (c :: TNET e).
(Ob a, Ob b, Ob c) =>
((a ** b) ** c) ~> (a ** (b ** c))
associator @(TN as) @(TN bs) @(TN cs) = forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) (cs :: [k]) r.
(KnownListOf c as, KnownListOf c bs, KnownListOf c cs) =>
((KnownListOf c ((as ++ bs) ++ cs),
  KnownListOf c (as ++ (bs ++ cs))) =>
 r)
-> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) (cs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs, KnownListOf c cs) =>
((KnownListOf c ((as ++ bs) ++ cs),
  KnownListOf c (as ++ (bs ++ cs))) =>
 r)
-> r
withKnownAssoc @KnownNat @as @bs @cs TensorNetwork
  (TN ((UN TN a ++ UN TN b) ++ UN TN c))
  (TN (UN TN a ++ (UN TN b ++ UN TN c)))
(KnownListOf KnownNat ((UN TN a ++ UN TN b) ++ UN TN c),
 KnownListOf KnownNat (UN TN a ++ (UN TN b ++ UN TN c))) =>
TensorNetwork
  (TN ((UN TN a ++ UN TN b) ++ UN TN c))
  (TN (UN TN a ++ (UN TN b ++ UN TN c)))
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
TensorNetwork (TN as) (TN bs)
straight
  associatorInv :: forall (a :: TNET e) (b :: TNET e) (c :: TNET e).
(Ob a, Ob b, Ob c) =>
(a ** (b ** c)) ~> ((a ** b) ** c)
associatorInv @(TN as) @(TN bs) @(TN cs) = forall {k} (c :: OB k) (as :: [k]) (bs :: [k]) (cs :: [k]) r.
(KnownListOf c as, KnownListOf c bs, KnownListOf c cs) =>
((KnownListOf c ((as ++ bs) ++ cs),
  KnownListOf c (as ++ (bs ++ cs))) =>
 r)
-> r
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]) (cs :: [Nat]) r.
(KnownListOf c as, KnownListOf c bs, KnownListOf c cs) =>
((KnownListOf c ((as ++ bs) ++ cs),
  KnownListOf c (as ++ (bs ++ cs))) =>
 r)
-> r
withKnownAssoc @KnownNat @as @bs @cs TensorNetwork
  (TN (UN TN a ++ (UN TN b ++ UN TN c)))
  (TN ((UN TN a ++ UN TN b) ++ UN TN c))
(KnownListOf KnownNat ((UN TN a ++ UN TN b) ++ UN TN c),
 KnownListOf KnownNat (UN TN a ++ (UN TN b ++ UN TN c))) =>
TensorNetwork
  (TN (UN TN a ++ (UN TN b ++ UN TN c)))
  (TN ((UN TN a ++ UN TN b) ++ UN TN c))
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
TensorNetwork (TN as) (TN bs)
straight

instance (Scalar e) => SymMonoidal (TNET e) where
  swap :: forall (a :: TNET e) (b :: TNET e).
(Ob a, Ob b) =>
(a ** b) ~> (b ** a)
swap @(TN as) @(TN bs) =
    forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @as @bs ((KnownListOf KnownNat (UN TN a ++ UN TN b) =>
  (a ** b) ~> (b ** a))
 -> (a ** b) ~> (b ** a))
-> (KnownListOf KnownNat (UN TN a ++ UN TN b) =>
    (a ** b) ~> (b ** a))
-> (a ** b) ~> (b ** a)
forall a b. (a -> b) -> a -> b
$
      forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @bs @as ((KnownListOf KnownNat (UN TN b ++ UN TN a) =>
  (a ** b) ~> (b ** a))
 -> (a ** b) ~> (b ** a))
-> (KnownListOf KnownNat (UN TN b ++ UN TN a) =>
    (a ** b) ~> (b ** a))
-> (a ** b) ~> (b ** a)
forall a b. (a -> b) -> a -> b
$
        let na :: Node
na = [Node] -> Node
forall a. [a] -> Node
forall (t :: Type -> Type) a. Foldable t => t a -> Node
P.length (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf @as)
            nb :: Node
nb = [Node] -> Node
forall a. [a] -> Node
forall (t :: Type -> Type) a. Foldable t => t a -> Node
P.length (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf @bs)
        in [Node]
-> [Node]
-> [Node]
-> TensorNetwork
     (TN (UN TN a ++ UN TN b)) (TN (UN TN b ++ UN TN a))
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
[Node] -> [Node] -> [Node] -> TensorNetwork (TN as) (TN bs)
wiring (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf @as [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf @bs) [Node
0 .. Node
na Node -> Node -> Node
forall a. Num a => a -> a -> a
P.+ Node
nb Node -> Node -> Node
forall a. Num a => a -> a -> a
- Node
1] ([Node
na .. Node
na Node -> Node -> Node
forall a. Num a => a -> a -> a
P.+ Node
nb Node -> Node -> Node
forall a. Num a => a -> a -> a
- Node
1] [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ [Node
0 .. Node
na Node -> Node -> Node
forall a. Num a => a -> a -> a
- Node
1])

-- | The distributors reorder the entries of the direct sums.
instance (Scalar e) => Distributive (TNET e) where
  distL :: forall (a :: TNET e) (b :: TNET e) (c :: TNET e).
(Ob a, Ob b, Ob c) =>
(a ** (b || c)) ~> ((a ** b) || (a ** c))
distL @(TN as) @(TN bs) @(TN cs) =
    forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownNat (Size as + Size bs) => r) -> r
withSum @bs @cs ((KnownNat (Size (UN TN b) + Size (UN TN c)) =>
  (a ** (b || c)) ~> ((a ** b) || (a ** c)))
 -> (a ** (b || c)) ~> ((a ** b) || (a ** c)))
-> (KnownNat (Size (UN TN b) + Size (UN TN c)) =>
    (a ** (b || c)) ~> ((a ** b) || (a ** c)))
-> (a ** (b || c)) ~> ((a ** b) || (a ** c))
forall a b. (a -> b) -> a -> b
$
      forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @as @'[Size bs + Size cs] ((KnownListOf
    KnownNat (UN TN a ++ '[Size (UN TN b) + Size (UN TN c)]) =>
  (a ** (b || c)) ~> ((a ** b) || (a ** c)))
 -> (a ** (b || c)) ~> ((a ** b) || (a ** c)))
-> (KnownListOf
      KnownNat (UN TN a ++ '[Size (UN TN b) + Size (UN TN c)]) =>
    (a ** (b || c)) ~> ((a ** b) || (a ** c)))
-> (a ** (b || c)) ~> ((a ** b) || (a ** c))
forall a b. (a -> b) -> a -> b
$
        forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @as @bs ((KnownListOf KnownNat (UN TN a ++ UN TN b) =>
  (a ** (b || c)) ~> ((a ** b) || (a ** c)))
 -> (a ** (b || c)) ~> ((a ** b) || (a ** c)))
-> (KnownListOf KnownNat (UN TN a ++ UN TN b) =>
    (a ** (b || c)) ~> ((a ** b) || (a ** c)))
-> (a ** (b || c)) ~> ((a ** b) || (a ** c))
forall a b. (a -> b) -> a -> b
$
          forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @as @cs ((KnownListOf KnownNat (UN TN a ++ UN TN c) =>
  (a ** (b || c)) ~> ((a ** b) || (a ** c)))
 -> (a ** (b || c)) ~> ((a ** b) || (a ** c)))
-> (KnownListOf KnownNat (UN TN a ++ UN TN c) =>
    (a ** (b || c)) ~> ((a ** b) || (a ** c)))
-> (a ** (b || c)) ~> ((a ** b) || (a ** c))
forall a b. (a -> b) -> a -> b
$
            forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownNat (Size as + Size bs) => r) -> r
withSum @(as ++ bs) @(as ++ cs) ((KnownNat
    (Size (UN TN a ++ UN TN b) + Size (UN TN a ++ UN TN c)) =>
  (a ** (b || c)) ~> ((a ** b) || (a ** c)))
 -> (a ** (b || c)) ~> ((a ** b) || (a ** c)))
-> (KnownNat
      (Size (UN TN a ++ UN TN b) + Size (UN TN a ++ UN TN c)) =>
    (a ** (b || c)) ~> ((a ** b) || (a ** c)))
-> (a ** (b || c)) ~> ((a ** b) || (a ** c))
forall a b. (a -> b) -> a -> b
$
              let (Node
sa, Node
sb, Node
sc) = (forall (ns :: [Nat]). KnownListOf KnownNat ns => Node
size @as, forall (ns :: [Nat]). KnownListOf KnownNat ns => Node
size @bs, forall (ns :: [Nat]). KnownListOf KnownNat ns => Node
size @cs)
                  target :: Node -> Node
target Node
c = let (Node
i, Node
k) = Node
c Node -> Node -> (Node, Node)
forall a. Integral a => a -> a -> (a, a)
`P.quotRem` (Node
sb Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
sc) in if Node
k Node -> Node -> Bool
forall a. Ord a => a -> a -> Bool
P.< Node
sb then Node
i Node -> Node -> Node
forall a. Num a => a -> a -> a
* Node
sb Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
k else Node
sa Node -> Node -> Node
forall a. Num a => a -> a -> a
* Node
sb Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
i Node -> Node -> Node
forall a. Num a => a -> a -> a
* Node
sc Node -> Node -> Node
forall a. Num a => a -> a -> a
+ Node
k Node -> Node -> Node
forall a. Num a => a -> a -> a
- Node
sb
              in (Node -> Node)
-> TensorNetwork
     (TN (UN TN a ++ '[Size (UN TN b) + Size (UN TN c)]))
     (TN '[Size (UN TN a ++ UN TN b) + Size (UN TN a ++ UN TN c)])
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Node -> Node) -> TensorNetwork (TN as) (TN bs)
reindex Node -> Node
target
  distR :: forall (a :: TNET e) (b :: TNET e) (c :: TNET e).
(Ob a, Ob b, Ob c) =>
((a || b) ** c) ~> ((a ** c) || (b ** c))
distR @(TN as) @(TN bs) @(TN cs) =
    forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownNat (Size as + Size bs) => r) -> r
withSum @as @bs ((KnownNat (Size (UN TN a) + Size (UN TN b)) =>
  ((a || b) ** c) ~> ((a ** c) || (b ** c)))
 -> ((a || b) ** c) ~> ((a ** c) || (b ** c)))
-> (KnownNat (Size (UN TN a) + Size (UN TN b)) =>
    ((a || b) ** c) ~> ((a ** c) || (b ** c)))
-> ((a || b) ** c) ~> ((a ** c) || (b ** c))
forall a b. (a -> b) -> a -> b
$
      forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @as @cs ((KnownListOf KnownNat (UN TN a ++ UN TN c) =>
  ((a || b) ** c) ~> ((a ** c) || (b ** c)))
 -> ((a || b) ** c) ~> ((a ** c) || (b ** c)))
-> (KnownListOf KnownNat (UN TN a ++ UN TN c) =>
    ((a || b) ** c) ~> ((a ** c) || (b ** c)))
-> ((a || b) ** c) ~> ((a ** c) || (b ** c))
forall a b. (a -> b) -> a -> b
$
        forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @bs @cs ((KnownListOf KnownNat (UN TN b ++ UN TN c) =>
  ((a || b) ** c) ~> ((a ** c) || (b ** c)))
 -> ((a || b) ** c) ~> ((a ** c) || (b ** c)))
-> (KnownListOf KnownNat (UN TN b ++ UN TN c) =>
    ((a || b) ** c) ~> ((a ** c) || (b ** c)))
-> ((a || b) ** c) ~> ((a ** c) || (b ** c))
forall a b. (a -> b) -> a -> b
$
          forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownNat (Size as + Size bs) => r) -> r
withSum @(as ++ cs) @(bs ++ cs) ((KnownNat
    (Size (UN TN a ++ UN TN c) + Size (UN TN b ++ UN TN c)) =>
  ((a || b) ** c) ~> ((a ** c) || (b ** c)))
 -> ((a || b) ** c) ~> ((a ** c) || (b ** c)))
-> (KnownNat
      (Size (UN TN a ++ UN TN c) + Size (UN TN b ++ UN TN c)) =>
    ((a || b) ** c) ~> ((a ** c) || (b ** c)))
-> ((a || b) ** c) ~> ((a ** c) || (b ** c))
forall a b. (a -> b) -> a -> b
$
            (Node -> Node)
-> TensorNetwork
     (TN ((Size (UN TN a) + Size (UN TN b)) : UN TN c))
     (TN '[Size (UN TN a ++ UN TN c) + Size (UN TN b ++ UN TN c)])
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(Node -> Node) -> TensorNetwork (TN as) (TN bs)
reindex Node -> Node
forall a. a -> a
P.id
  absorbL :: forall (a :: TNET e). Ob a => (a ** InitialObject) ~> InitialObject
absorbL @(TN as) = forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @as @'[0] TensorNetwork (TN (UN TN a ++ '[0])) (TN '[0])
KnownListOf KnownNat (UN TN a ++ '[0]) =>
TensorNetwork (TN (UN TN a ++ '[0])) (TN '[0])
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
TensorNetwork (TN as) (TN bs)
zero
  absorbR :: forall (a :: TNET e). Ob a => (InitialObject ** a) ~> InitialObject
absorbR = (InitialObject ** a) ~> InitialObject
TensorNetwork (TN (0 : UN TN a)) (TN '[0])
forall {e} (as :: [Nat]) (bs :: [Nat]).
(Scalar e, KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
TensorNetwork (TN as) (TN bs)
zero

-- | Every object is self-dual, with the transpose as the dual of an arrow: its wiring the other way
-- round.
instance (Scalar e) => Dialogue (TNET e) where
  type Dual a = a
  withObDual :: forall (a :: TNET e) r. Ob a => (Ob (Dual a) => r) -> r
withObDual Ob (Dual a) => r
r = r
Ob (Dual a) => r
r
  dual :: forall (a :: TNET e) (b :: TNET e). (a ~> b) -> Dual b ~> Dual a
dual = (a ~> b) -> Dual b ~> Dual a
TensorNetwork a b -> TensorNetwork b a
forall e (a :: TNET e) (b :: TNET e).
TensorNetwork a b -> TensorNetwork b a
transpose
  linDist :: forall (a :: TNET e) (b :: TNET e) (c :: TNET e).
(Ob a, Ob b, Ob c) =>
((a ** b) ~> Dual c) -> a ~> Dual (b ** c)
linDist @(TN as) @(TN bs) @(TN cs) (TensorNetwork (Net [Node]
d [Node]
i [Node]
o [Factor e]
f)) =
    forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @bs @cs ((KnownListOf KnownNat (UN TN b ++ UN TN c) => a ~> Dual (b ** c))
 -> a ~> Dual (b ** c))
-> (KnownListOf KnownNat (UN TN b ++ UN TN c) =>
    a ~> Dual (b ** c))
-> a ~> Dual (b ** c)
forall a b. (a -> b) -> a -> b
$ let na :: Node
na = [Node] -> Node
forall a. [a] -> Node
forall (t :: Type -> Type) a. Foldable t => t a -> Node
P.length (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf @as) in Net e -> TensorNetwork (TN (UN TN a)) (TN (UN TN b ++ UN TN c))
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Net e -> TensorNetwork (TN as) (TN bs)
TensorNetwork ([Node] -> [Node] -> [Node] -> [Factor e] -> Net e
forall e. [Node] -> [Node] -> [Node] -> [Factor e] -> Net e
Net [Node]
d (Node -> [Node] -> [Node]
forall a. Node -> [a] -> [a]
P.take Node
na [Node]
i) (Node -> [Node] -> [Node]
forall a. Node -> [a] -> [a]
P.drop Node
na [Node]
i [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ [Node]
o) [Factor e]
f)
  linDistInv :: forall (a :: TNET e) (b :: TNET e) (c :: TNET e).
(Ob a, Ob b, Ob c) =>
(a ~> Dual (b ** c)) -> (a ** b) ~> Dual c
linDistInv @(TN as) @(TN bs) (TensorNetwork (Net [Node]
d [Node]
i [Node]
o [Factor e]
f)) =
    forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @as @bs ((KnownListOf KnownNat (UN TN a ++ UN TN b) => (a ** b) ~> Dual c)
 -> (a ** b) ~> Dual c)
-> (KnownListOf KnownNat (UN TN a ++ UN TN b) =>
    (a ** b) ~> Dual c)
-> (a ** b) ~> Dual c
forall a b. (a -> b) -> a -> b
$ let nb :: Node
nb = [Node] -> Node
forall a. [a] -> Node
forall (t :: Type -> Type) a. Foldable t => t a -> Node
P.length (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf @bs) in Net e -> TensorNetwork (TN (UN TN a ++ UN TN b)) (TN (UN TN c))
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
Net e -> TensorNetwork (TN as) (TN bs)
TensorNetwork ([Node] -> [Node] -> [Node] -> [Factor e] -> Net e
forall e. [Node] -> [Node] -> [Node] -> [Factor e] -> Net e
Net [Node]
d ([Node]
i [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ Node -> [Node] -> [Node]
forall a. Node -> [a] -> [a]
P.take Node
nb [Node]
o) (Node -> [Node] -> [Node]
forall a. Node -> [a] -> [a]
P.drop Node
nb [Node]
o) [Factor e]
f)
  doubleNegInv :: forall (a :: TNET e). Ob a => a ~> Dual (Dual a)
doubleNegInv = a ~> Dual (Dual a)
TensorNetwork (TN (UN TN a)) (TN (UN TN a))
forall {k} (p :: CAT k) (a :: k). (Promonad p, Ob a) => p a a
forall (a :: TNET e). Ob a => TensorNetwork a a
id

instance (Scalar e) => StarAutonomous (TNET e) where
  dualInv :: forall (a :: TNET e) (b :: TNET e).
(Ob a, Ob b) =>
(Dual a ~> Dual b) -> b ~> a
dualInv = (Dual a ~> Dual b) -> b ~> a
TensorNetwork (TN (UN TN a)) (TN (UN TN b))
-> TensorNetwork (TN (UN TN b)) (TN (UN TN a))
forall e (a :: TNET e) (b :: TNET e).
TensorNetwork a b -> TensorNetwork b a
transpose
  doubleNeg :: forall (a :: TNET e). Ob a => Dual (Dual a) ~> a
doubleNeg = Dual (Dual a) ~> a
TensorNetwork (TN (UN TN a)) (TN (UN TN a))
forall {k} (p :: CAT k) (a :: k). (Promonad p, Ob a) => p a a
forall (a :: TNET e). Ob a => TensorNetwork a a
id

instance (Scalar e) => Closed (TNET e) where
  type a ~~> b = ExpSA a b
  withObExp :: forall (a :: TNET e) (b :: TNET e) r.
(Ob a, Ob b) =>
(Ob (a ~~> b) => r) -> r
withObExp @(TN as) @(TN bs) Ob (a ~~> b) => r
r = forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @as @bs r
Ob (a ~~> b) => r
KnownListOf KnownNat (UN TN a ++ UN TN b) => r
r
  curry :: forall (a :: TNET e) (b :: TNET e) (c :: TNET e).
(Ob a, Ob b) =>
((a ** b) ~> c) -> a ~> (b ~~> c)
curry @a @b = forall {k} (a :: k) (b :: k) (c :: k).
(StarAutonomous k, Ob a, Ob b) =>
((a ** b) ~> c) -> a ~> ExpSA b c
forall (a :: TNET e) (b :: TNET e) (c :: TNET e).
(StarAutonomous (TNET e), Ob a, Ob b) =>
((a ** b) ~> c) -> a ~> ExpSA b c
currySA @a @b
  apply :: forall (a :: TNET e) (b :: TNET e).
(Ob a, Ob b) =>
((a ~~> b) ** a) ~> b
apply @a @b = forall {k} (b :: k) (c :: k).
(StarAutonomous k, Ob b, Ob c) =>
(ExpSA b c ** b) ~> c
forall (b :: TNET e) (c :: TNET e).
(StarAutonomous (TNET e), Ob b, Ob c) =>
(ExpSA b c ** b) ~> c
applySA @a @b
  ^^^ :: forall (a :: TNET e) (b :: TNET e) (x :: TNET e) (y :: TNET e).
(b ~> y) -> (x ~> a) -> (a ~~> b) ~> (x ~~> y)
(^^^) = (b ~> y) -> (x ~> a) -> ExpSA a b ~> ExpSA x y
(b ~> y) -> (x ~> a) -> (a ~~> b) ~> (x ~~> y)
forall {k} (a :: k) (b :: k) (x :: k) (y :: k).
StarAutonomous k =>
(b ~> y) -> (x ~> a) -> ExpSA a b ~> ExpSA x y
expSA

instance (Scalar e) => IsoMix (TNET e) where
  dualUnit :: Dual Unit ~> Unit
dualUnit = Dual Unit ~> Unit
TensorNetwork (TN '[]) (TN '[])
forall {k} (p :: CAT k) (a :: k). (Promonad p, Ob a) => p a a
forall (a :: TNET e). Ob a => TensorNetwork a a
id
  dualUnitInv :: Unit ~> Dual Unit
dualUnitInv = Unit ~> Dual Unit
TensorNetwork (TN '[]) (TN '[])
forall {k} (p :: CAT k) (a :: k). (Promonad p, Ob a) => p a a
forall (a :: TNET e). Ob a => TensorNetwork a a
id
  dualityCounit :: forall (a :: TNET e). Ob a => (Dual a ** a) ~> Unit
dualityCounit @a = forall {k} (a :: k). Frobenius a => (a ** a) ~> Unit
forall (a :: TNET e). Frobenius a => (a ** a) ~> Unit
cap @a

instance (Scalar e) => CompactClosed (TNET e) where
  distribDual :: forall (a :: TNET e) (b :: TNET e).
(Ob a, Ob b) =>
Dual (a ** b) ~> (Dual a ** Dual b)
distribDual @(TN as) @(TN bs) = forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @as @bs TensorNetwork (TN (UN TN a ++ UN TN b)) (TN (UN TN a ++ UN TN b))
KnownListOf KnownNat (UN TN a ++ UN TN b) =>
TensorNetwork (TN (UN TN a ++ UN TN b)) (TN (UN TN a ++ UN TN b))
forall {k} (p :: CAT k) (a :: k). (Promonad p, Ob a) => p a a
forall (a :: TNET e). Ob a => TensorNetwork a a
id
  dualityUnit :: forall (a :: TNET e). Ob a => Unit ~> (a ** Dual a)
dualityUnit @a = forall {k} (a :: k). Frobenius a => Unit ~> (a ** a)
forall (a :: TNET e). Frobenius a => Unit ~> (a ** a)
cup @a

instance (Scalar e, MonoidalAction (t :: (TNET e, TNET e) +-> TNET e)) => Costrong t (TensorNetwork :: CAT (TNET e)) where
  coact :: forall (a :: TNET e) (x :: TNET e) (y :: TNET e).
(Ob a, Ob x, Ob y) =>
TensorNetwork (Act t a x) (Act t a y) -> TensorNetwork x y
coact @x = forall {m} {k} (t :: (m, k) +-> k) (u :: m) (x :: k) (y :: k).
(CompactClosed m, MonoidalAction t, Ob x, Ob y, Ob u) =>
(Act t u x ~> Act t u y) -> x ~> y
forall (t :: (TNET e, TNET e) +-> TNET e) (u :: TNET e)
       (x :: TNET e) (y :: TNET e).
(CompactClosed (TNET e), MonoidalAction t, Ob x, Ob y, Ob u) =>
(Act t u x ~> Act t u y) -> x ~> y
coactCC @t @x

-- | Merging each wire with its partner, and the unit.
instance (Scalar e, KnownListOf KnownNat ns) => Monoid (TN ns :: TNET e) where
  mempty :: Unit ~> TN ns
mempty = [Node] -> [Node] -> [Node] -> TensorNetwork (TN '[]) (TN ns)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
[Node] -> [Node] -> [Node] -> TensorNetwork (TN as) (TN bs)
wiring (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf @ns) [] (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
wires @ns)
  mappend :: (TN ns ** TN ns) ~> TN ns
mappend = forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @ns @ns ([Node] -> [Node] -> [Node] -> TensorNetwork (TN (ns ++ ns)) (TN ns)
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
[Node] -> [Node] -> [Node] -> TensorNetwork (TN as) (TN bs)
wiring (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf @ns) (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
wires @ns [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
wires @ns) (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
wires @ns))

-- | Copying each wire, and discarding.
instance (Scalar e, KnownListOf KnownNat ns) => Comonoid (TN ns :: TNET e) where
  counit :: TN ns ~> Unit
counit = [Node] -> [Node] -> [Node] -> TensorNetwork (TN ns) (TN '[])
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
[Node] -> [Node] -> [Node] -> TensorNetwork (TN as) (TN bs)
wiring (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf @ns) (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
wires @ns) []
  comult :: TN ns ~> (TN ns ** TN ns)
comult = forall (as :: [Nat]) (bs :: [Nat]) r.
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
(KnownListOf KnownNat (as ++ bs) => r) -> r
withAppend @ns @ns ([Node] -> [Node] -> [Node] -> TensorNetwork (TN ns) (TN (ns ++ ns))
forall {e} (as :: [Nat]) (bs :: [Nat]).
(KnownListOf KnownNat as, KnownListOf KnownNat bs) =>
[Node] -> [Node] -> [Node] -> TensorNetwork (TN as) (TN bs)
wiring (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf @ns) (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
wires @ns) (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
wires @ns [Node] -> [Node] -> [Node]
forall a. [a] -> [a] -> [a]
P.++ forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
wires @ns))

-- | A node for each wire of a list.
wires :: forall ns. (KnownListOf KnownNat ns) => [Node]
wires :: forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
wires = [Node
0 .. [Node] -> Node
forall a. [a] -> Node
forall (t :: Type -> Type) a. Foldable t => t a -> Node
P.length (forall (ns :: [Nat]). KnownListOf KnownNat ns => [Node]
dimsOf @ns) Node -> Node -> Node
forall a. Num a => a -> a -> a
- Node
1]

instance (Scalar e, KnownListOf KnownNat ns) => CommutativeMonoid (TN ns :: TNET e)
instance (Scalar e, KnownListOf KnownNat ns) => CocommutativeComonoid (TN ns :: TNET e)
instance (Scalar e, KnownListOf KnownNat ns) => Frobenius (TN ns :: TNET e)
instance (Scalar e) => Hypergraph (TNET e)
instance (Scalar e) => CopyDiscard (TNET e)

-- | The size of an object is the product of its dimensions.
instance (Scalar e) => Sized (TNET e) where
  sizeOf :: forall (a :: TNET e). Ob a => Node
sizeOf @(TN ns) = forall (ns :: [Nat]). KnownListOf KnownNat ns => Node
size @ns

-- | Two objects are isomorphic when they have the same dimensions.
instance (Scalar e) => DecidableIso (TNET e) where
  isoOf :: forall (c :: (TNET e +-> TNET e) -> Constraint) (a :: TNET e)
       (b :: TNET e).
(Ob a, Ob b) =>
Maybe (Optic c a a b b)
isoOf @_ @(TN as) @(TN bs) =
    Maybe (a :~: b) -> Maybe (Optic c a a b b)
forall {k} (c :: (k +-> k) -> Constraint) (a :: k) (b :: k).
(CategoryOf k, Ob a) =>
Maybe (a :~: b) -> Maybe (Optic c a a b b)
isoFromEquality
      ( ((UN TN a :~: UN TN b) -> a :~: b)
-> Maybe (UN TN a :~: UN TN b) -> Maybe (a :~: b)
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
P.fmap
          (\UN TN a :~: UN TN b
Refl -> a :~: a
a :~: b
forall {k} (a :: k). a :~: a
Refl)
          (forall {k} (c :: OB k) (as :: [k]) (bs :: [k]).
(forall (x :: k) (y :: k). (c x, c y) => Maybe (x :~: y))
-> ListOf c as -> ListOf c bs -> Maybe (as :~: bs)
forall (c :: OB Nat) (as :: [Nat]) (bs :: [Nat]).
(forall (x :: Nat) (y :: Nat). (c x, c y) => Maybe (x :~: y))
-> ListOf c as -> ListOf c bs -> Maybe (as :~: bs)
eqListOf @KnownNat (\ @x @y -> Proxy x -> Proxy y -> Maybe (x :~: y)
forall (a :: Nat) (b :: Nat) (proxy1 :: Nat -> Type)
       (proxy2 :: Nat -> Type).
(KnownNat a, KnownNat b) =>
proxy1 a -> proxy2 b -> Maybe (a :~: b)
sameNat (forall (t :: Nat). Proxy t
forall {k} (t :: k). Proxy t
Proxy @x) (forall (t :: Nat). Proxy t
forall {k} (t :: k). Proxy t
Proxy @y)) (forall {k} (c :: OB k) (xs :: [k]). KnownListOf c xs => ListOf c xs
forall (c :: OB Nat) (xs :: [Nat]). KnownListOf c xs => ListOf c xs
listOf @KnownNat @as) (forall {k} (c :: OB k) (xs :: [k]). KnownListOf c xs => ListOf c xs
forall (c :: OB Nat) (xs :: [Nat]). KnownListOf c xs => ListOf c xs
listOf @KnownNat @bs))
      )