-- Copyright 2020 Google LLC
--
-- Licensed under the Apache License, Version 2.0 (the "License");
-- you may not use this file except in compliance with the License.
-- You may obtain a copy of the License at
--
--      http://www.apache.org/licenses/LICENSE-2.0
--
-- Unless required by applicable law or agreed to in writing, software
-- distributed under the License is distributed on an "AS IS" BASIS,
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-- See the License for the specific language governing permissions and
-- limitations under the License.

{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
-- | Arrays of static size.  The arrays are polymorphic in the underlying
-- linear data structure used to store the actual values.
module Data.Array.Internal.ShapedG(
  Array(..), Shape(..), Size, Rank, Vector, VecElem,
  Window, Stride, Permute, Permutation, ValidDims,
  Broadcast,
  size, shapeL, rank,
  toList, fromList, toVector, fromVector,
  normalize,
  scalar, unScalar, constant,
  reshape, stretch, stretchOuter, transpose,
  index, pad,
  mapA, zipWithA, zipWith3A,
  append,
  ravel, unravel,
  window, stride,
  slice, rerank, rerank2, rev,
  reduce, foldrA, traverseA,
  allSameA,
  sumA, productA, minimumA, maximumA,
  anyA, allA,
  broadcast,
  generate, iterateN, iota,
  ) where
import Control.DeepSeq
import Data.Data(Data)
import Data.Proxy(Proxy(..))
import GHC.Generics(Generic)
import GHC.Stack(HasCallStack)
import GHC.TypeLits(Nat, type (<=), KnownNat, type (+))
import Test.QuickCheck hiding (generate)
import Text.PrettyPrint.HughesPJClass

import Data.Array.Internal
import Data.Array.Internal.Shape

-- | Arrays stored in a /v/ with values of type /a/.
newtype Array (sh :: [Nat]) v a = A (T v a)
  deriving ((forall x. Array sh v a -> Rep (Array sh v a) x)
-> (forall x. Rep (Array sh v a) x -> Array sh v a)
-> Generic (Array sh v a)
forall (sh :: [Nat]) (v :: * -> *) a x.
Rep (Array sh v a) x -> Array sh v a
forall (sh :: [Nat]) (v :: * -> *) a x.
Array sh v a -> Rep (Array sh v a) x
forall x. Rep (Array sh v a) x -> Array sh v a
forall x. Array sh v a -> Rep (Array sh v a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall (sh :: [Nat]) (v :: * -> *) a x.
Rep (Array sh v a) x -> Array sh v a
$cfrom :: forall (sh :: [Nat]) (v :: * -> *) a x.
Array sh v a -> Rep (Array sh v a) x
Generic, Typeable (Array sh v a)
DataType
Constr
Typeable (Array sh v a)
-> (forall (c :: * -> *).
    (forall d b. Data d => c (d -> b) -> d -> c b)
    -> (forall g. g -> c g) -> Array sh v a -> c (Array sh v a))
-> (forall (c :: * -> *).
    (forall b r. Data b => c (b -> r) -> c r)
    -> (forall r. r -> c r) -> Constr -> c (Array sh v a))
-> (Array sh v a -> Constr)
-> (Array sh v a -> DataType)
-> (forall (t :: * -> *) (c :: * -> *).
    Typeable t =>
    (forall d. Data d => c (t d)) -> Maybe (c (Array sh v a)))
-> (forall (t :: * -> * -> *) (c :: * -> *).
    Typeable t =>
    (forall d e. (Data d, Data e) => c (t d e))
    -> Maybe (c (Array sh v a)))
-> ((forall b. Data b => b -> b) -> Array sh v a -> Array sh v a)
-> (forall r r'.
    (r -> r' -> r)
    -> r -> (forall d. Data d => d -> r') -> Array sh v a -> r)
-> (forall r r'.
    (r' -> r -> r)
    -> r -> (forall d. Data d => d -> r') -> Array sh v a -> r)
-> (forall u. (forall d. Data d => d -> u) -> Array sh v a -> [u])
-> (forall u.
    Int -> (forall d. Data d => d -> u) -> Array sh v a -> u)
-> (forall (m :: * -> *).
    Monad m =>
    (forall d. Data d => d -> m d) -> Array sh v a -> m (Array sh v a))
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> Array sh v a -> m (Array sh v a))
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> Array sh v a -> m (Array sh v a))
-> Data (Array sh v a)
Array sh v a -> DataType
Array sh v a -> Constr
(forall b. Data b => b -> b) -> Array sh v a -> Array sh v a
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Array sh v a -> c (Array sh v a)
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Array sh v a)
forall (sh :: [Nat]) (v :: * -> *) a.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
Typeable (Array sh v a)
forall (sh :: [Nat]) (v :: * -> *) a.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
Array sh v a -> DataType
forall (sh :: [Nat]) (v :: * -> *) a.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
Array sh v a -> Constr
forall (sh :: [Nat]) (v :: * -> *) a.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
(forall b. Data b => b -> b) -> Array sh v a -> Array sh v a
forall (sh :: [Nat]) (v :: * -> *) a u.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
Int -> (forall d. Data d => d -> u) -> Array sh v a -> u
forall (sh :: [Nat]) (v :: * -> *) a u.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
(forall d. Data d => d -> u) -> Array sh v a -> [u]
forall (sh :: [Nat]) (v :: * -> *) a r r'.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
(r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh v a -> r
forall (sh :: [Nat]) (v :: * -> *) a r r'.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
(r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh v a -> r
forall (sh :: [Nat]) (v :: * -> *) a (m :: * -> *).
(Typeable v, Typeable sh, Typeable a, Data (v a), Monad m) =>
(forall d. Data d => d -> m d) -> Array sh v a -> m (Array sh v a)
forall (sh :: [Nat]) (v :: * -> *) a (m :: * -> *).
(Typeable v, Typeable sh, Typeable a, Data (v a), MonadPlus m) =>
(forall d. Data d => d -> m d) -> Array sh v a -> m (Array sh v a)
forall (sh :: [Nat]) (v :: * -> *) a (c :: * -> *).
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Array sh v a)
forall (sh :: [Nat]) (v :: * -> *) a (c :: * -> *).
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Array sh v a -> c (Array sh v a)
forall (sh :: [Nat]) (v :: * -> *) a (t :: * -> *) (c :: * -> *).
(Typeable v, Typeable sh, Typeable a, Data (v a), Typeable t) =>
(forall d. Data d => c (t d)) -> Maybe (c (Array sh v a))
forall (sh :: [Nat]) (v :: * -> *) a (t :: * -> * -> *)
       (c :: * -> *).
(Typeable v, Typeable sh, Typeable a, Data (v a), Typeable t) =>
(forall d e. (Data d, Data e) => c (t d e))
-> Maybe (c (Array sh v a))
forall a.
Typeable a
-> (forall (c :: * -> *).
    (forall d b. Data d => c (d -> b) -> d -> c b)
    -> (forall g. g -> c g) -> a -> c a)
-> (forall (c :: * -> *).
    (forall b r. Data b => c (b -> r) -> c r)
    -> (forall r. r -> c r) -> Constr -> c a)
-> (a -> Constr)
-> (a -> DataType)
-> (forall (t :: * -> *) (c :: * -> *).
    Typeable t =>
    (forall d. Data d => c (t d)) -> Maybe (c a))
-> (forall (t :: * -> * -> *) (c :: * -> *).
    Typeable t =>
    (forall d e. (Data d, Data e) => c (t d e)) -> Maybe (c a))
-> ((forall b. Data b => b -> b) -> a -> a)
-> (forall r r'.
    (r -> r' -> r) -> r -> (forall d. Data d => d -> r') -> a -> r)
-> (forall r r'.
    (r' -> r -> r) -> r -> (forall d. Data d => d -> r') -> a -> r)
-> (forall u. (forall d. Data d => d -> u) -> a -> [u])
-> (forall u. Int -> (forall d. Data d => d -> u) -> a -> u)
-> (forall (m :: * -> *).
    Monad m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> a -> m a)
-> Data a
forall u. Int -> (forall d. Data d => d -> u) -> Array sh v a -> u
forall u. (forall d. Data d => d -> u) -> Array sh v a -> [u]
forall r r'.
(r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh v a -> r
forall r r'.
(r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh v a -> r
forall (m :: * -> *).
Monad m =>
(forall d. Data d => d -> m d) -> Array sh v a -> m (Array sh v a)
forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d) -> Array sh v a -> m (Array sh v a)
forall (c :: * -> *).
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Array sh v a)
forall (c :: * -> *).
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Array sh v a -> c (Array sh v a)
forall (t :: * -> *) (c :: * -> *).
Typeable t =>
(forall d. Data d => c (t d)) -> Maybe (c (Array sh v a))
forall (t :: * -> * -> *) (c :: * -> *).
Typeable t =>
(forall d e. (Data d, Data e) => c (t d e))
-> Maybe (c (Array sh v a))
$cA :: Constr
$tArray :: DataType
gmapMo :: (forall d. Data d => d -> m d) -> Array sh v a -> m (Array sh v a)
$cgmapMo :: forall (sh :: [Nat]) (v :: * -> *) a (m :: * -> *).
(Typeable v, Typeable sh, Typeable a, Data (v a), MonadPlus m) =>
(forall d. Data d => d -> m d) -> Array sh v a -> m (Array sh v a)
gmapMp :: (forall d. Data d => d -> m d) -> Array sh v a -> m (Array sh v a)
$cgmapMp :: forall (sh :: [Nat]) (v :: * -> *) a (m :: * -> *).
(Typeable v, Typeable sh, Typeable a, Data (v a), MonadPlus m) =>
(forall d. Data d => d -> m d) -> Array sh v a -> m (Array sh v a)
gmapM :: (forall d. Data d => d -> m d) -> Array sh v a -> m (Array sh v a)
$cgmapM :: forall (sh :: [Nat]) (v :: * -> *) a (m :: * -> *).
(Typeable v, Typeable sh, Typeable a, Data (v a), Monad m) =>
(forall d. Data d => d -> m d) -> Array sh v a -> m (Array sh v a)
gmapQi :: Int -> (forall d. Data d => d -> u) -> Array sh v a -> u
$cgmapQi :: forall (sh :: [Nat]) (v :: * -> *) a u.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
Int -> (forall d. Data d => d -> u) -> Array sh v a -> u
gmapQ :: (forall d. Data d => d -> u) -> Array sh v a -> [u]
$cgmapQ :: forall (sh :: [Nat]) (v :: * -> *) a u.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
(forall d. Data d => d -> u) -> Array sh v a -> [u]
gmapQr :: (r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh v a -> r
$cgmapQr :: forall (sh :: [Nat]) (v :: * -> *) a r r'.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
(r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh v a -> r
gmapQl :: (r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh v a -> r
$cgmapQl :: forall (sh :: [Nat]) (v :: * -> *) a r r'.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
(r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh v a -> r
gmapT :: (forall b. Data b => b -> b) -> Array sh v a -> Array sh v a
$cgmapT :: forall (sh :: [Nat]) (v :: * -> *) a.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
(forall b. Data b => b -> b) -> Array sh v a -> Array sh v a
dataCast2 :: (forall d e. (Data d, Data e) => c (t d e))
-> Maybe (c (Array sh v a))
$cdataCast2 :: forall (sh :: [Nat]) (v :: * -> *) a (t :: * -> * -> *)
       (c :: * -> *).
(Typeable v, Typeable sh, Typeable a, Data (v a), Typeable t) =>
(forall d e. (Data d, Data e) => c (t d e))
-> Maybe (c (Array sh v a))
dataCast1 :: (forall d. Data d => c (t d)) -> Maybe (c (Array sh v a))
$cdataCast1 :: forall (sh :: [Nat]) (v :: * -> *) a (t :: * -> *) (c :: * -> *).
(Typeable v, Typeable sh, Typeable a, Data (v a), Typeable t) =>
(forall d. Data d => c (t d)) -> Maybe (c (Array sh v a))
dataTypeOf :: Array sh v a -> DataType
$cdataTypeOf :: forall (sh :: [Nat]) (v :: * -> *) a.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
Array sh v a -> DataType
toConstr :: Array sh v a -> Constr
$ctoConstr :: forall (sh :: [Nat]) (v :: * -> *) a.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
Array sh v a -> Constr
gunfold :: (forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Array sh v a)
$cgunfold :: forall (sh :: [Nat]) (v :: * -> *) a (c :: * -> *).
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Array sh v a)
gfoldl :: (forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Array sh v a -> c (Array sh v a)
$cgfoldl :: forall (sh :: [Nat]) (v :: * -> *) a (c :: * -> *).
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Array sh v a -> c (Array sh v a)
$cp1Data :: forall (sh :: [Nat]) (v :: * -> *) a.
(Typeable v, Typeable sh, Typeable a, Data (v a)) =>
Typeable (Array sh v a)
Data)

instance (Vector v, Show a, VecElem v a, Shape sh, Show (v a)) => Show (Array sh v a) where
  showsPrec :: Int -> Array sh v a -> ShowS
showsPrec Int
p a :: Array sh v a
a@(A T v a
_) = Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10) (ShowS -> ShowS) -> ShowS -> ShowS
forall a b. (a -> b) -> a -> b
$
    String -> ShowS
showString String
"fromList @" ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> ShapeL -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
11 (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShowS
showStringString
" " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [a] -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
11 (Array sh v a -> [a]
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Shape sh) =>
Array sh v a -> [a]
toList Array sh v a
a)

instance (Shape sh, Vector v, Read a, VecElem v a) => Read (Array sh v a) where
  readsPrec :: Int -> ReadS (Array sh v a)
readsPrec Int
p = Bool -> ReadS (Array sh v a) -> ReadS (Array sh v a)
forall a. Bool -> ReadS a -> ReadS a
readParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
10) (ReadS (Array sh v a) -> ReadS (Array sh v a))
-> ReadS (Array sh v a) -> ReadS (Array sh v a)
forall a b. (a -> b) -> a -> b
$ \ String
r1 ->
    [([a] -> Array sh v a
forall (sh :: [Nat]) (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a, Shape sh) =>
[a] -> Array sh v a
fromList [a]
xs, String
r4)
    | (String
"fromList", String
r2) <- ReadS String
lex String
r1, (String
"@", String
r2') <- ReadS String
lex String
r2
    , (ShapeL
s, String
r3) <- Int -> ReadS ShapeL
forall a. Read a => Int -> ReadS a
readsPrec Int
11 String
r2', ([a]
xs, String
r4) <- Int -> ReadS [a]
forall a. Read a => Int -> ReadS a
readsPrec Int
11 String
r3
    , ShapeL
s ShapeL -> ShapeL -> Bool
forall a. Eq a => a -> a -> Bool
== Proxy sh -> ShapeL
forall (s :: [Nat]). Shape s => Proxy s -> ShapeL
shapeP (Proxy sh
forall k (t :: k). Proxy t
Proxy :: Proxy sh), ShapeL -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ShapeL
s Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs]

instance (Vector v, Eq a, VecElem v a, Eq (v a), Shape sh)
         => Eq (Array sh v a) where
  a :: Array sh v a
a@(A T v a
v) == :: Array sh v a -> Array sh v a -> Bool
== (A T v a
v') = ShapeL -> T v a -> T v a -> Bool
forall (v :: * -> *) a.
(Vector v, VecElem v a, Eq a, Eq (v a)) =>
ShapeL -> T v a -> T v a -> Bool
equalT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) T v a
v T v a
v'
  {-# INLINE (==) #-}

instance (Vector v, Ord a, Ord (v a), VecElem v a, Shape sh)
         => Ord (Array sh v a) where
  a :: Array sh v a
a@(A T v a
v) compare :: Array sh v a -> Array sh v a -> Ordering
`compare` (A T v a
v') = ShapeL -> T v a -> T v a -> Ordering
forall (v :: * -> *) a.
(Vector v, VecElem v a, Ord a, Ord (v a)) =>
ShapeL -> T v a -> T v a -> Ordering
compareT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) T v a
v T v a
v'
  {-# INLINE compare #-}

instance (Vector v, Pretty a, VecElem v a, Shape sh) => Pretty (Array sh v a) where
  pPrintPrec :: PrettyLevel -> Rational -> Array sh v a -> Doc
pPrintPrec PrettyLevel
l Rational
p a :: Array sh v a
a@(A T v a
t) = PrettyLevel -> Rational -> ShapeL -> T v a -> Doc
forall (v :: * -> *) a.
(Vector v, VecElem v a, Pretty a) =>
PrettyLevel -> Rational -> ShapeL -> T v a -> Doc
ppT PrettyLevel
l Rational
p (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) T v a
t

instance (NFData (v a)) => NFData (Array sh v a) where
  rnf :: Array sh v a -> ()
rnf (A T v a
t) = T v a -> ()
forall a. NFData a => a -> ()
rnf T v a
t

-- | The number of elements in the array.
{-# INLINE size #-}
size :: forall sh v a . (Shape sh) => Array sh v a -> Int
size :: Array sh v a -> Int
size Array sh v a
_ = Proxy sh -> Int
forall (s :: [Nat]). Shape s => Proxy s -> Int
sizeP (Proxy sh
forall k (t :: k). Proxy t
Proxy :: Proxy sh)

-- | The shape of an array, i.e., a list of the sizes of its dimensions.
-- In the linearization of the array the outermost (i.e. first list element)
-- varies most slowly.
-- O(1) time.
{-# INLINE shapeL #-}
shapeL :: forall sh v a . (Shape sh) => Array sh v a -> ShapeL
shapeL :: Array sh v a -> ShapeL
shapeL Array sh v a
_ = Proxy sh -> ShapeL
forall (s :: [Nat]). Shape s => Proxy s -> ShapeL
shapeP (Proxy sh
forall k (t :: k). Proxy t
Proxy :: Proxy sh)

-- | The rank of an array, i.e., the number if dimensions it has.
-- O(1) time.
{-# INLINE rank #-}
rank :: forall sh v a . (Shape sh, KnownNat (Rank sh)) => Array sh v a -> Int
rank :: Array sh v a -> Int
rank Array sh v a
_ = forall i. (KnownNat (Rank sh), Num i) => i
forall (n :: Nat) i. (KnownNat n, Num i) => i
valueOf @(Rank sh)

-- | Index into an array.  Fails if the array has rank 0 or if the index is out of bounds.
-- O(1) time.
{-# INLINE index #-}
index :: forall s sh v a . (HasCallStack, Vector v, KnownNat s) =>
         Array (s:sh) v a -> Int -> Array sh v a
index :: Array (s : sh) v a -> Int -> Array sh v a
index (A T v a
t) Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
s = String -> Array sh v a
forall a. HasCallStack => String -> a
error (String -> Array sh v a) -> String -> Array sh v a
forall a b. (a -> b) -> a -> b
$ String
"index: out of bounds " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
i String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" >= " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
s
              | Bool
otherwise = T v a -> Array sh v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v a -> Array sh v a) -> T v a -> Array sh v a
forall a b. (a -> b) -> a -> b
$ T v a -> Int -> T v a
forall (v :: * -> *) a. T v a -> Int -> T v a
indexT T v a
t Int
i
  where s :: Int
s = forall i. (KnownNat s, Num i) => i
forall (n :: Nat) i. (KnownNat n, Num i) => i
valueOf @s

-- | Convert to a list with the elements in the linearization order.
-- O(1) time.
{-# INLINE toList #-}
toList :: (Vector v, VecElem v a, Shape sh) => Array sh v a -> [a]
toList :: Array sh v a -> [a]
toList a :: Array sh v a
a@(A T v a
t) = ShapeL -> T v a -> [a]
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
ShapeL -> T v a -> [a]
toListT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) T v a
t

-- | Convert to a vector with the elements in the linearization order.
-- O(n) or O(1) time (the latter if the vector is already in the linearization order).
{-# INLINE toVector #-}
toVector :: (Vector v, VecElem v a, Shape sh) => Array sh v a -> v a
toVector :: Array sh v a -> v a
toVector a :: Array sh v a
a@(A T v a
t) = ShapeL -> T v a -> v a
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
ShapeL -> T v a -> v a
toVectorT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) T v a
t

-- | Convert from a list with the elements given in the linearization order.
-- Fails if the given shape does not have the same number of elements as the list.
-- O(n) time.
{-# INLINE fromList #-}
fromList :: forall sh v a . (HasCallStack, Vector v, VecElem v a, Shape sh) =>
            [a] -> Array sh v a
fromList :: [a] -> Array sh v a
fromList [a]
vs | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
l = String -> Array sh v a
forall a. HasCallStack => String -> a
error (String -> Array sh v a) -> String -> Array sh v a
forall a b. (a -> b) -> a -> b
$ String
"fromList: size mismatch " String -> ShowS
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> String
forall a. Show a => a -> String
show (Int
n, Int
l)
            | Bool
otherwise = T v a -> Array sh v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v a -> Array sh v a) -> T v a -> Array sh v a
forall a b. (a -> b) -> a -> b
$ ShapeL -> Int -> v a -> T v a
forall (v :: * -> *) a. ShapeL -> Int -> v a -> T v a
T ShapeL
st Int
0 (v a -> T v a) -> v a -> T v a
forall a b. (a -> b) -> a -> b
$ [a] -> v a
forall (v :: * -> *) a. (Vector v, VecElem v a) => [a] -> v a
vFromList [a]
vs
  where Int
n : ShapeL
st = ShapeL -> ShapeL
getStridesT ShapeL
ss
        l :: Int
l = [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
vs
        ss :: ShapeL
ss = Proxy sh -> ShapeL
forall (s :: [Nat]). Shape s => Proxy s -> ShapeL
shapeP (Proxy sh
forall k (t :: k). Proxy t
Proxy :: Proxy sh)

-- | Convert from a vector with the elements given in the linearization order.
-- Fails if the given shape does not have the same number of elements as the list.
-- O(1) time.
{-# INLINE fromVector #-}
fromVector :: forall sh v a . (HasCallStack, Vector v, VecElem v a, Shape sh) =>
              v a -> Array sh v a
fromVector :: v a -> Array sh v a
fromVector v a
v | Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
l = String -> Array sh v a
forall a. HasCallStack => String -> a
error (String -> Array sh v a) -> String -> Array sh v a
forall a b. (a -> b) -> a -> b
$ String
"fromVector: size mismatch" String -> ShowS
forall a. [a] -> [a] -> [a]
++ (Int, Int) -> String
forall a. Show a => a -> String
show (Int
n, Int
l)
             | Bool
otherwise = T v a -> Array sh v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v a -> Array sh v a) -> T v a -> Array sh v a
forall a b. (a -> b) -> a -> b
$ ShapeL -> Int -> v a -> T v a
forall (v :: * -> *) a. ShapeL -> Int -> v a -> T v a
T ShapeL
st Int
0 v a
v
  where Int
n : ShapeL
st = ShapeL -> ShapeL
getStridesT ShapeL
ss
        l :: Int
l = v a -> Int
forall (v :: * -> *) a. (Vector v, VecElem v a) => v a -> Int
vLength v a
v
        ss :: ShapeL
ss = Proxy sh -> ShapeL
forall (s :: [Nat]). Shape s => Proxy s -> ShapeL
shapeP (Proxy sh
forall k (t :: k). Proxy t
Proxy :: Proxy sh)

-- | Make sure the underlying vector is in the linearization order.
-- This is semantically an identity function, but can have big performance
-- implications.
-- O(n) or O(1) time.
{-# INLINE normalize #-}
normalize :: (Vector v, VecElem v a, Shape sh) => Array sh v a -> Array sh v a
normalize :: Array sh v a -> Array sh v a
normalize = v a -> Array sh v a
forall (sh :: [Nat]) (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a, Shape sh) =>
v a -> Array sh v a
fromVector (v a -> Array sh v a)
-> (Array sh v a -> v a) -> Array sh v a -> Array sh v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh v a -> v a
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Shape sh) =>
Array sh v a -> v a
toVector

-- | Change the shape of an array.  Type error if the arrays have different number of elements.
-- O(n) or O(1) time.
{-# INLINE reshape #-}
reshape :: forall sh' sh v a .
           (Vector v, VecElem v a, Shape sh, Shape sh', Size sh ~ Size sh') =>
           Array sh v a -> Array sh' v a
reshape :: Array sh v a -> Array sh' v a
reshape Array sh v a
a = ShapeL -> ShapeL -> Array sh v a -> Array sh' v a
forall (v :: * -> *) a (sh :: [Nat]) (sh' :: [Nat]).
(Vector v, VecElem v a) =>
ShapeL -> ShapeL -> Array sh v a -> Array sh' v a
reshape' (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) (Proxy sh' -> ShapeL
forall (s :: [Nat]). Shape s => Proxy s -> ShapeL
shapeP (Proxy sh'
forall k (t :: k). Proxy t
Proxy :: Proxy sh')) Array sh v a
a

reshape' :: (Vector v, VecElem v a) =>
            ShapeL -> ShapeL -> Array sh v a -> Array sh' v a
reshape' :: ShapeL -> ShapeL -> Array sh v a -> Array sh' v a
reshape' ShapeL
sh ShapeL
sh' (A t :: T v a
t@(T ShapeL
ost Int
oo v a
v))
  | v a -> Int
forall (v :: * -> *) a. (Vector v, VecElem v a) => v a -> Int
vLength v a
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = T v a -> Array sh' v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v a -> Array sh' v a) -> T v a -> Array sh' v a
forall a b. (a -> b) -> a -> b
$ ShapeL -> Int -> v a -> T v a
forall (v :: * -> *) a. ShapeL -> Int -> v a -> T v a
T ((Int -> Int) -> ShapeL -> ShapeL
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Int -> Int
forall a b. a -> b -> a
const Int
0) ShapeL
sh) Int
0 v a
v  -- Fast special case for singleton vector
  | Just ShapeL
nst <- ShapeL -> ShapeL -> ShapeL -> Maybe ShapeL
simpleReshape ShapeL
ost ShapeL
sh ShapeL
sh' = T v a -> Array sh' v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v a -> Array sh' v a) -> T v a -> Array sh' v a
forall a b. (a -> b) -> a -> b
$ ShapeL -> Int -> v a -> T v a
forall (v :: * -> *) a. ShapeL -> Int -> v a -> T v a
T ShapeL
nst Int
oo v a
v
  | Bool
otherwise = T v a -> Array sh' v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v a -> Array sh' v a) -> T v a -> Array sh' v a
forall a b. (a -> b) -> a -> b
$ ShapeL -> v a -> T v a
forall (v :: * -> *) a. ShapeL -> v a -> T v a
fromVectorT ShapeL
sh' (v a -> T v a) -> v a -> T v a
forall a b. (a -> b) -> a -> b
$ ShapeL -> T v a -> v a
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
ShapeL -> T v a -> v a
toVectorT ShapeL
sh T v a
t

-- | Change the size of dimensions with size 1.  These dimension can be changed to any size.
-- All other dimensions must remain the same.
-- O(1) time.
{-# INLINE stretch #-}
stretch :: forall sh' sh v a . (Shape sh, Shape sh', ValidStretch sh sh') =>
           Array sh v a -> Array sh' v a
stretch :: Array sh v a -> Array sh' v a
stretch = [Bool] -> Array sh v a -> Array sh' v a
forall (sh :: [Nat]) (v :: * -> *) a (sh' :: [Nat]).
[Bool] -> Array sh v a -> Array sh' v a
stretch' (Proxy sh -> Proxy sh' -> [Bool]
forall (from :: [Nat]) (to :: [Nat]).
ValidStretch from to =>
Proxy from -> Proxy to -> [Bool]
stretching (Proxy sh
forall k (t :: k). Proxy t
Proxy :: Proxy sh) (Proxy sh'
forall k (t :: k). Proxy t
Proxy :: Proxy sh'))

stretch' :: [Bool] -> Array sh v a -> Array sh' v a
stretch' :: [Bool] -> Array sh v a -> Array sh' v a
stretch' [Bool]
str (A T v a
vs) = T v a -> Array sh' v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v a -> Array sh' v a) -> T v a -> Array sh' v a
forall a b. (a -> b) -> a -> b
$ [Bool] -> T v a -> T v a
forall (v :: * -> *) a. [Bool] -> T v a -> T v a
stretchT [Bool]
str T v a
vs

-- | Change the size of the outermost dimension by replication.
{-# INLINE stretchOuter #-}
stretchOuter :: forall s sh v a . (Shape sh) =>
                Array (1 : sh) v a -> Array (s : sh) v a
stretchOuter :: Array (1 : sh) v a -> Array (s : sh) v a
stretchOuter (A T v a
vs) = T v a -> Array (s : sh) v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v a -> Array (s : sh) v a) -> T v a -> Array (s : sh) v a
forall a b. (a -> b) -> a -> b
$ [Bool] -> T v a -> T v a
forall (v :: * -> *) a. [Bool] -> T v a -> T v a
stretchT (Bool
True Bool -> [Bool] -> [Bool]
forall a. a -> [a] -> [a]
: (Int -> Bool) -> ShapeL -> [Bool]
forall a b. (a -> b) -> [a] -> [b]
map (Bool -> Int -> Bool
forall a b. a -> b -> a
const Bool
False) (T v a -> ShapeL
forall (v :: * -> *) a. T v a -> ShapeL
strides T v a
vs)) T v a
vs

-- | Convert a value to a scalar (rank 0) array.
-- O(1) time.
{-# INLINE scalar #-}
scalar :: (Vector v, VecElem v a) => a -> Array '[] v a
scalar :: a -> Array '[] v a
scalar = T v a -> Array '[] v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v a -> Array '[] v a) -> (a -> T v a) -> a -> Array '[] v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> T v a
forall (v :: * -> *) a. (Vector v, VecElem v a) => a -> T v a
scalarT

-- | Convert a scalar (rank 0) array to a value.
-- O(1) time.
{-# INLINE unScalar #-}
unScalar :: (Vector v, VecElem v a) => Array '[] v a -> a
unScalar :: Array '[] v a -> a
unScalar (A T v a
t) = T v a -> a
forall (v :: * -> *) a. (Vector v, VecElem v a) => T v a -> a
unScalarT T v a
t

-- | Make an array with all elements having the same value.
-- O(1) time.
{-# INLINE constant #-}
constant :: forall sh v a . (Vector v, VecElem v a, Shape sh) =>
            a -> Array sh v a
constant :: a -> Array sh v a
constant = T v a -> Array sh v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v a -> Array sh v a) -> (a -> T v a) -> a -> Array sh v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> a -> T v a
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
ShapeL -> a -> T v a
constantT (Proxy sh -> ShapeL
forall (s :: [Nat]). Shape s => Proxy s -> ShapeL
shapeP (Proxy sh
forall k (t :: k). Proxy t
Proxy :: Proxy sh))

-- | Map over the array elements.
-- O(n) time.
{-# INLINE mapA #-}
mapA :: (Vector v, VecElem v a, VecElem v b, Shape sh) =>
        (a -> b) -> Array sh v a -> Array sh v b
mapA :: (a -> b) -> Array sh v a -> Array sh v b
mapA a -> b
f a :: Array sh v a
a@(A T v a
t) = T v b -> Array sh v b
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v b -> Array sh v b) -> T v b -> Array sh v b
forall a b. (a -> b) -> a -> b
$ ShapeL -> (a -> b) -> T v a -> T v b
forall (v :: * -> *) a b.
(Vector v, VecElem v a, VecElem v b) =>
ShapeL -> (a -> b) -> T v a -> T v b
mapT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) a -> b
f T v a
t

-- | Map over the array elements.
-- O(n) time.
{-# INLINE zipWithA #-}
zipWithA :: (Vector v, VecElem v a, VecElem v b, VecElem v c, Shape sh) =>
            (a -> b -> c) -> Array sh v a -> Array sh v b -> Array sh v c
zipWithA :: (a -> b -> c) -> Array sh v a -> Array sh v b -> Array sh v c
zipWithA a -> b -> c
f a :: Array sh v a
a@(A T v a
t) (A T v b
t') = T v c -> Array sh v c
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v c -> Array sh v c) -> T v c -> Array sh v c
forall a b. (a -> b) -> a -> b
$ ShapeL -> (a -> b -> c) -> T v a -> T v b -> T v c
forall (v :: * -> *) a b c.
(Vector v, VecElem v a, VecElem v b, VecElem v c) =>
ShapeL -> (a -> b -> c) -> T v a -> T v b -> T v c
zipWithT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) a -> b -> c
f T v a
t T v b
t'

-- | Map over the array elements.
-- O(n) time.
{-# INLINE zipWith3A #-}
zipWith3A :: (Vector v, VecElem v a, VecElem v b, VecElem v c, VecElem v d, Shape sh) =>
             (a -> b -> c -> d) -> Array sh v a -> Array sh v b -> Array sh v c -> Array sh v d
zipWith3A :: (a -> b -> c -> d)
-> Array sh v a -> Array sh v b -> Array sh v c -> Array sh v d
zipWith3A a -> b -> c -> d
f a :: Array sh v a
a@(A T v a
t) (A T v b
t') (A T v c
t'') = T v d -> Array sh v d
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v d -> Array sh v d) -> T v d -> Array sh v d
forall a b. (a -> b) -> a -> b
$ ShapeL -> (a -> b -> c -> d) -> T v a -> T v b -> T v c -> T v d
forall (v :: * -> *) a b c d.
(Vector v, VecElem v a, VecElem v b, VecElem v c, VecElem v d) =>
ShapeL -> (a -> b -> c -> d) -> T v a -> T v b -> T v c -> T v d
zipWith3T (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) a -> b -> c -> d
f T v a
t T v b
t' T v c
t''

-- | Pad each dimension on the low and high side with the given value.
-- O(n) time.
{-# INLINE pad #-}
pad :: forall ps sh' sh a v . (HasCallStack, Vector v, VecElem v a, Padded ps sh sh', Shape sh) =>
       a -> Array sh v a -> Array sh' v a
pad :: a -> Array sh v a -> Array sh' v a
pad a
v a :: Array sh v a
a@(A T v a
at) = T v a -> Array sh' v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v a -> Array sh' v a) -> T v a -> Array sh' v a
forall a b. (a -> b) -> a -> b
$ (ShapeL, T v a) -> T v a
forall a b. (a, b) -> b
snd ((ShapeL, T v a) -> T v a) -> (ShapeL, T v a) -> T v a
forall a b. (a -> b) -> a -> b
$ a -> [(Int, Int)] -> ShapeL -> T v a -> (ShapeL, T v a)
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
a -> [(Int, Int)] -> ShapeL -> T v a -> (ShapeL, T v a)
padT a
v [(Int, Int)]
aps ShapeL
ash T v a
at
  where ash :: ShapeL
ash = Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a
        aps :: [(Int, Int)]
aps = Proxy ps -> Proxy sh -> [(Int, Int)]
forall (ps :: [(Nat, Nat)]) (sh :: [Nat]) (sh' :: [Nat]).
Padded ps sh sh' =>
Proxy ps -> Proxy sh -> [(Int, Int)]
padded (Proxy ps
forall k (t :: k). Proxy t
Proxy :: Proxy ps) (Proxy sh
forall k (t :: k). Proxy t
Proxy :: Proxy sh)

-- | Do an arbitrary array transposition.
-- Fails if the transposition argument is not a permutation of the numbers
-- [0..r-1], where r is the rank of the array.
-- O(1) time.
{-# INLINE transpose #-}
transpose :: forall is sh v a .
             (Permutation is, Rank is <= Rank sh, Shape sh, Shape is, KnownNat (Rank sh)) =>
             Array sh v a -> Array (Permute is sh) v a
transpose :: Array sh v a -> Array (Permute is sh) v a
transpose (A T v a
t) = T v a -> Array (Permute is sh) v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (ShapeL -> T v a -> T v a
forall (v :: * -> *) a. ShapeL -> T v a -> T v a
transposeT ShapeL
is' T v a
t)
  where l :: Int
l = ShapeL -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ShapeL
is
        n :: Int
n = forall i. (KnownNat (Rank sh), Num i) => i
forall (n :: Nat) i. (KnownNat n, Num i) => i
valueOf @(Rank sh)
        is' :: ShapeL
is' = ShapeL
is ShapeL -> ShapeL -> ShapeL
forall a. [a] -> [a] -> [a]
++ [Int
l .. Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1]
        is :: ShapeL
is = Proxy is -> ShapeL
forall (s :: [Nat]). Shape s => Proxy s -> ShapeL
shapeP (Proxy is
forall k (t :: k). Proxy t
Proxy :: Proxy is)

-- | Append two arrays along the outermost dimension.
-- All dimensions, except the outermost, must be the same.
-- O(n) time.
{-# INLINE append #-}
append :: (Vector v, VecElem v a, Shape sh, KnownNat m, KnownNat n, KnownNat (m+n)) =>
          Array (m ': sh) v a -> Array (n ': sh) v a -> Array (m+n ': sh) v a
append :: Array (m : sh) v a
-> Array (n : sh) v a -> Array ((m + n) : sh) v a
append Array (m : sh) v a
a Array (n : sh) v a
b = v a -> Array ((m + n) : sh) v a
forall (sh :: [Nat]) (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a, Shape sh) =>
v a -> Array sh v a
fromVector (v a -> v a -> v a
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
v a -> v a -> v a
vAppend (Array (m : sh) v a -> v a
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Shape sh) =>
Array sh v a -> v a
toVector Array (m : sh) v a
a) (Array (n : sh) v a -> v a
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Shape sh) =>
Array sh v a -> v a
toVector Array (n : sh) v a
b))

-- | Turn a rank-1 array of arrays into a single array by making the outer array into the outermost
-- dimension of the result array.  All the arrays must have the same shape.
-- O(n) time.
{-# INLINE ravel #-}
ravel :: (Vector v, Vector v', VecElem v a, VecElem v' (Array sh v a)
         , Shape sh, KnownNat s) =>
         Array '[s] v' (Array sh v a) -> Array (s:sh) v a
ravel :: Array '[s] v' (Array sh v a) -> Array (s : sh) v a
ravel = v a -> Array (s : sh) v a
forall (sh :: [Nat]) (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a, Shape sh) =>
v a -> Array sh v a
fromVector (v a -> Array (s : sh) v a)
-> (Array '[s] v' (Array sh v a) -> v a)
-> Array '[s] v' (Array sh v a)
-> Array (s : sh) v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [v a] -> v a
forall (v :: * -> *) a. (Vector v, VecElem v a) => [v a] -> v a
vConcat ([v a] -> v a)
-> (Array '[s] v' (Array sh v a) -> [v a])
-> Array '[s] v' (Array sh v a)
-> v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array sh v a -> v a) -> [Array sh v a] -> [v a]
forall a b. (a -> b) -> [a] -> [b]
map Array sh v a -> v a
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Shape sh) =>
Array sh v a -> v a
toVector ([Array sh v a] -> [v a])
-> (Array '[s] v' (Array sh v a) -> [Array sh v a])
-> Array '[s] v' (Array sh v a)
-> [v a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array '[s] v' (Array sh v a) -> [Array sh v a]
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Shape sh) =>
Array sh v a -> [a]
toList

-- | Turn an array into a nested array, this is the inverse of 'ravel'.
-- I.e., @ravel . unravel == id@.
-- O(n) time.
{-# INLINE unravel #-}
unravel :: (Vector v, Vector v', VecElem v a, VecElem v' (Array sh v a)
           , Shape sh, KnownNat s) =>
           Array (s:sh) v a -> Array '[s] v' (Array sh v a)
unravel :: Array (s : sh) v a -> Array '[s] v' (Array sh v a)
unravel = (Array sh v a -> Array '[] v' (Array sh v a))
-> Array (s : sh) v a
-> Array (Take 1 (s : sh) ++ '[]) v' (Array sh v a)
forall (n :: Nat) (i :: [Nat]) (o :: [Nat]) (sh :: [Nat])
       (v :: * -> *) (v' :: * -> *) a b.
(Vector v, Vector v', VecElem v a, VecElem v' b, Drop n sh ~ i,
 Shape sh, KnownNat n, Shape o, Shape (Take n sh ++ o)) =>
(Array i v a -> Array o v' b)
-> Array sh v a -> Array (Take n sh ++ o) v' b
rerank @1 Array sh v a -> Array '[] v' (Array sh v a)
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
a -> Array '[] v a
scalar

-- | Make a window of the outermost dimensions.
-- The rank increases with the length of the window list.
-- E.g., if the shape of the array is @[10,12,8]@ and
-- the window size is @[3,3]@ then the resulting array will have shape
-- @[8,10,3,3,8]@.
--
-- E.g., @window [2] (fromList [4] [1,2,3,4]) == fromList [3,2] [1,2, 2,3, 3,4]@
-- O(1) time.
--
-- If the window parameter @ws = [w1,...,wk]@ and @wa = window ws a@ then
-- @wa `index` i1 ... `index` ik == slice [(i1,w1),...,(ik,wk)] a@.
{-# INLINE window #-}
window :: forall ws sh' sh v a .
          (Window ws sh sh', Vector v, KnownNat (Rank ws)) =>
          Array sh v a -> Array sh' v a
window :: Array sh v a -> Array sh' v a
window (A (T ShapeL
ss Int
o v a
v)) = T v a -> Array sh' v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (ShapeL -> Int -> v a -> T v a
forall (v :: * -> *) a. ShapeL -> Int -> v a -> T v a
T (ShapeL
ss' ShapeL -> ShapeL -> ShapeL
forall a. [a] -> [a] -> [a]
++ ShapeL
ss) Int
o v a
v)
  where ss' :: ShapeL
ss' = Int -> ShapeL -> ShapeL
forall a. Int -> [a] -> [a]
take (forall i. (KnownNat (Rank ws), Num i) => i
forall (n :: Nat) i. (KnownNat n, Num i) => i
valueOf @(Rank ws)) ShapeL
ss

-- | Stride the outermost dimensions.
-- E.g., if the array shape is @[10,12,8]@ and the strides are
-- @[2,2]@ then the resulting shape will be @[5,6,8]@.
-- O(1) time.
{-# INLINE stride #-}
stride :: forall ts sh' sh v a .
          (Stride ts sh sh', Vector v, Shape ts) =>
          Array sh v a -> Array sh' v a
stride :: Array sh v a -> Array sh' v a
stride (A (T ShapeL
ss Int
o v a
v)) = T v a -> Array sh' v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (ShapeL -> Int -> v a -> T v a
forall (v :: * -> *) a. ShapeL -> Int -> v a -> T v a
T ((Int -> Int -> Int) -> ShapeL -> ShapeL -> ShapeL
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) (ShapeL
ats ShapeL -> ShapeL -> ShapeL
forall a. [a] -> [a] -> [a]
++ Int -> ShapeL
forall a. a -> [a]
repeat Int
1) ShapeL
ss) Int
o v a
v)
  where ats :: ShapeL
ats = Proxy ts -> ShapeL
forall (s :: [Nat]). Shape s => Proxy s -> ShapeL
shapeP (Proxy ts
forall k (t :: k). Proxy t
Proxy :: Proxy ts)

-- | Extract a slice of an array.
-- The first type argument is a list of (offset, length) pairs.
-- The length of the slicing argument must not exceed the rank of the array.
-- The extracted slice must fall within the array dimensions.
-- E.g. @slice @'[ '(1,2)] (fromList @'[4] [1,2,3,4]) == fromList @'[2] [2,3]@.
-- O(1) time.
{-# INLINE slice #-}
slice :: forall sl sh' sh v a .
         (Slice sl sh sh') =>
         Array sh v a -> Array sh' v a
slice :: Array sh v a -> Array sh' v a
slice (A (T ShapeL
ts Int
o v a
v)) = T v a -> Array sh' v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (ShapeL -> Int -> v a -> T v a
forall (v :: * -> *) a. ShapeL -> Int -> v a -> T v a
T ShapeL
ts (Int
oInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
i) v a
v)
  where i :: Int
i = ShapeL -> Int
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (ShapeL -> Int) -> ShapeL -> Int
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> ShapeL -> ShapeL -> ShapeL
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) ShapeL
ts (ShapeL -> ShapeL) -> ShapeL -> ShapeL
forall a b. (a -> b) -> a -> b
$ Proxy sl -> Proxy sh -> ShapeL
forall (ls :: [(Nat, Nat)]) (ss :: [Nat]) (rs :: [Nat]).
Slice ls ss rs =>
Proxy ls -> Proxy ss -> ShapeL
sliceOffsets (Proxy sl
forall k (t :: k). Proxy t
Proxy :: Proxy sl) (Proxy sh
forall k (t :: k). Proxy t
Proxy :: Proxy sh)

-- | Apply a function to the subarrays /n/ levels down and make
-- the results into an array with the same /n/ outermost dimensions.
-- The /n/ must not exceed the rank of the array.
-- O(n) time.
{-# INLINE rerank #-}
rerank :: forall n i o sh v v' a b .
          (Vector v, Vector v', VecElem v a, VecElem v' b,
           Drop n sh ~ i, Shape sh, KnownNat n, Shape o, Shape (Take n sh ++ o)) =>
          (Array i v a -> Array o v' b) -> Array sh v a -> Array (Take n sh ++ o) v' b
rerank :: (Array i v a -> Array o v' b)
-> Array sh v a -> Array (Take n sh ++ o) v' b
rerank Array i v a -> Array o v' b
f a :: Array sh v a
a@(A T v a
t) =
  v' b -> Array (Take n sh ++ o) v' b
forall (sh :: [Nat]) (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a, Shape sh) =>
v a -> Array sh v a
fromVector (v' b -> Array (Take n sh ++ o) v' b)
-> v' b -> Array (Take n sh ++ o) v' b
forall a b. (a -> b) -> a -> b
$
  [v' b] -> v' b
forall (v :: * -> *) a. (Vector v, VecElem v a) => [v a] -> v a
vConcat ([v' b] -> v' b) -> [v' b] -> v' b
forall a b. (a -> b) -> a -> b
$
  (T v a -> v' b) -> [T v a] -> [v' b]
forall a b. (a -> b) -> [a] -> [b]
map (Array o v' b -> v' b
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Shape sh) =>
Array sh v a -> v a
toVector (Array o v' b -> v' b) -> (T v a -> Array o v' b) -> T v a -> v' b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array i v a -> Array o v' b
f (Array i v a -> Array o v' b)
-> (T v a -> Array i v a) -> T v a -> Array o v' b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T v a -> Array i v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A) ([T v a] -> [v' b]) -> [T v a] -> [v' b]
forall a b. (a -> b) -> a -> b
$
  ShapeL -> T v a -> [T v a]
forall (v :: * -> *) a. ShapeL -> T v a -> [T v a]
subArraysT ShapeL
osh T v a
t
  where osh :: ShapeL
osh = Int -> ShapeL -> ShapeL
forall a. Int -> [a] -> [a]
take (forall i. (KnownNat n, Num i) => i
forall (n :: Nat) i. (KnownNat n, Num i) => i
valueOf @n) (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a)

-- | Apply a two-argument function to the subarrays /n/ levels down and make
-- the results into an array with the same /n/ outermost dimensions.
-- The /n/ must not exceed the rank of the array.
-- O(n) time.
{-# INLINE rerank2 #-}
rerank2 :: forall n i1 i2 o sh1 sh2 r v a b c .
           (Vector v, VecElem v a, VecElem v b, VecElem v c,
            Drop n sh1 ~ i1, Drop n sh2 ~ i2, Shape sh1, Shape sh2,
            Take n sh1 ~ r, Take n sh2 ~ r,
            KnownNat n, Shape o, Shape (r ++ o)) =>
           (Array i1 v a -> Array i2 v b -> Array o v c) -> Array sh1 v a -> Array sh2 v b -> Array (r ++ o) v c
rerank2 :: (Array i1 v a -> Array i2 v b -> Array o v c)
-> Array sh1 v a -> Array sh2 v b -> Array (r ++ o) v c
rerank2 Array i1 v a -> Array i2 v b -> Array o v c
f aa :: Array sh1 v a
aa@(A T v a
ta) (A T v b
tb) =
  v c -> Array (r ++ o) v c
forall (sh :: [Nat]) (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a, Shape sh) =>
v a -> Array sh v a
fromVector (v c -> Array (r ++ o) v c) -> v c -> Array (r ++ o) v c
forall a b. (a -> b) -> a -> b
$
  [v c] -> v c
forall (v :: * -> *) a. (Vector v, VecElem v a) => [v a] -> v a
vConcat ([v c] -> v c) -> [v c] -> v c
forall a b. (a -> b) -> a -> b
$
  (T v a -> T v b -> v c) -> [T v a] -> [T v b] -> [v c]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\ T v a
a T v b
b -> Array o v c -> v c
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Shape sh) =>
Array sh v a -> v a
toVector (Array o v c -> v c) -> Array o v c -> v c
forall a b. (a -> b) -> a -> b
$ Array i1 v a -> Array i2 v b -> Array o v c
f (T v a -> Array i1 v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A T v a
a) (T v b -> Array i2 v b
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A T v b
b))
          (ShapeL -> T v a -> [T v a]
forall (v :: * -> *) a. ShapeL -> T v a -> [T v a]
subArraysT ShapeL
osh T v a
ta)
          (ShapeL -> T v b -> [T v b]
forall (v :: * -> *) a. ShapeL -> T v a -> [T v a]
subArraysT ShapeL
osh T v b
tb)
  where osh :: ShapeL
osh = Int -> ShapeL -> ShapeL
forall a. Int -> [a] -> [a]
take (forall i. (KnownNat n, Num i) => i
forall (n :: Nat) i. (KnownNat n, Num i) => i
valueOf @n) (Array sh1 v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh1 v a
aa)


-- | Reverse the given dimensions, with the outermost being dimension 0.
-- O(1) time.
{-# INLINE rev #-}
rev :: forall rs sh v a . (ValidDims rs sh, Shape rs, Shape sh) => Array sh v a -> Array sh v a
rev :: Array sh v a -> Array sh v a
rev a :: Array sh v a
a@(A T v a
t) = T v a -> Array sh v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (ShapeL -> ShapeL -> T v a -> T v a
forall (v :: * -> *) a. ShapeL -> ShapeL -> T v a -> T v a
reverseT ShapeL
rs ShapeL
sh T v a
t)
  where rs :: ShapeL
rs = Proxy rs -> ShapeL
forall (s :: [Nat]). Shape s => Proxy s -> ShapeL
shapeP (Proxy rs
forall k (t :: k). Proxy t
Proxy :: Proxy rs)
        sh :: ShapeL
sh = Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a

-- | Reduce all elements of an array into a rank 0 array.
-- To reduce parts use 'rerank' and 'transpose' together with 'reduce'.
-- O(n) time.
{-# INLINE reduce #-}
reduce :: (Vector v, VecElem v a, Shape sh) =>
          (a -> a -> a) -> a -> Array sh v a -> Array '[] v a
reduce :: (a -> a -> a) -> a -> Array sh v a -> Array '[] v a
reduce a -> a -> a
f a
z a :: Array sh v a
a@(A T v a
t) = T v a -> Array '[] v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v a -> Array '[] v a) -> T v a -> Array '[] v a
forall a b. (a -> b) -> a -> b
$ ShapeL -> (a -> a -> a) -> a -> T v a -> T v a
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
ShapeL -> (a -> a -> a) -> a -> T v a -> T v a
reduceT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) a -> a -> a
f a
z T v a
t

-- | Right fold across all elements of an array.
{-# INLINE foldrA #-}
foldrA
  :: (Vector v, VecElem v a, Shape sh)
  => (a -> b -> b) -> b -> Array sh v a -> b
foldrA :: (a -> b -> b) -> b -> Array sh v a -> b
foldrA a -> b -> b
f b
z a :: Array sh v a
a@(A T v a
t) = ShapeL -> (a -> b -> b) -> b -> T v a -> b
forall (v :: * -> *) a b.
(Vector v, VecElem v a) =>
ShapeL -> (a -> b -> b) -> b -> T v a -> b
foldrT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) a -> b -> b
f b
z T v a
t

-- | Constrained version of 'traverse' for 'Array's.
{-# INLINE traverseA #-}
traverseA
  :: (Vector v, VecElem v a, VecElem v b, Applicative f, Shape sh)
  => (a -> f b) -> Array sh v a -> f (Array sh v b)
traverseA :: (a -> f b) -> Array sh v a -> f (Array sh v b)
traverseA a -> f b
f a :: Array sh v a
a@(A T v a
t) = T v b -> Array sh v b
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v b -> Array sh v b) -> f (T v b) -> f (Array sh v b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ShapeL -> (a -> f b) -> T v a -> f (T v b)
forall (v :: * -> *) a b (f :: * -> *).
(Vector v, VecElem v a, VecElem v b, Applicative f) =>
ShapeL -> (a -> f b) -> T v a -> f (T v b)
traverseT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) a -> f b
f T v a
t

-- | Check if all elements of the array are equal.
allSameA :: (Shape sh, Vector v, VecElem v a, Eq a) => Array sh v a -> Bool
allSameA :: Array sh v a -> Bool
allSameA a :: Array sh v a
a@(A T v a
t) = ShapeL -> T v a -> Bool
forall (v :: * -> *) a.
(Vector v, VecElem v a, Eq a) =>
ShapeL -> T v a -> Bool
allSameT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) T v a
t

instance (Shape sh, Vector v, VecElem v a, Arbitrary a) => Arbitrary (Array sh v a) where
  arbitrary :: Gen (Array sh v a)
arbitrary = [a] -> Array sh v a
forall (sh :: [Nat]) (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a, Shape sh) =>
[a] -> Array sh v a
fromList ([a] -> Array sh v a) -> Gen [a] -> Gen (Array sh v a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Gen [a]
forall a. Arbitrary a => Int -> Gen [a]
vector (Proxy sh -> Int
forall (s :: [Nat]). Shape s => Proxy s -> Int
sizeP (Proxy sh
forall k (t :: k). Proxy t
Proxy :: Proxy sh))

-- | Sum of all elements.
{-# INLINE sumA #-}
sumA :: (Vector v, VecElem v a, Num a, Shape sh) => Array sh v a -> a
sumA :: Array sh v a -> a
sumA a :: Array sh v a
a@(A T v a
t) = ShapeL -> T v a -> a
forall (v :: * -> *) a.
(Vector v, VecElem v a, Num a) =>
ShapeL -> T v a -> a
sumT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) T v a
t

-- | Product of all elements.
{-# INLINE productA #-}
productA :: (Vector v, VecElem v a, Num a, Shape sh) => Array sh v a -> a
productA :: Array sh v a -> a
productA a :: Array sh v a
a@(A T v a
t) = ShapeL -> T v a -> a
forall (v :: * -> *) a.
(Vector v, VecElem v a, Num a) =>
ShapeL -> T v a -> a
productT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) T v a
t

-- | Maximum of all elements.
{-# INLINE maximumA #-}
maximumA :: (Vector v, VecElem v a, Ord a, Shape sh, 1 <= Size sh) => Array sh v a -> a
maximumA :: Array sh v a -> a
maximumA a :: Array sh v a
a@(A T v a
t) = ShapeL -> T v a -> a
forall (v :: * -> *) a.
(Vector v, VecElem v a, Ord a) =>
ShapeL -> T v a -> a
maximumT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) T v a
t

-- | Minimum of all elements.
{-# INLINE minimumA #-}
minimumA :: (Vector v, VecElem v a, Ord a, Shape sh, 1 <= Size sh) => Array sh v a -> a
minimumA :: Array sh v a -> a
minimumA a :: Array sh v a
a@(A T v a
t) = ShapeL -> T v a -> a
forall (v :: * -> *) a.
(Vector v, VecElem v a, Ord a) =>
ShapeL -> T v a -> a
minimumT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) T v a
t

-- | Test if the predicate holds for any element.
{-# INLINE anyA #-}
anyA :: (Vector v, VecElem v a, Shape sh) => (a -> Bool) -> Array sh v a -> Bool
anyA :: (a -> Bool) -> Array sh v a -> Bool
anyA a -> Bool
p a :: Array sh v a
a@(A T v a
t) = ShapeL -> (a -> Bool) -> T v a -> Bool
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
ShapeL -> (a -> Bool) -> T v a -> Bool
anyT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) a -> Bool
p T v a
t

-- | Test if the predicate holds for all elements.
{-# INLINE allA #-}
allA :: (Vector v, VecElem v a, Shape sh) => (a -> Bool) -> Array sh v a -> Bool
allA :: (a -> Bool) -> Array sh v a -> Bool
allA a -> Bool
p a :: Array sh v a
a@(A T v a
t) = ShapeL -> (a -> Bool) -> T v a -> Bool
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
ShapeL -> (a -> Bool) -> T v a -> Bool
anyT (Array sh v a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
shapeL Array sh v a
a) a -> Bool
p T v a
t

-- | Put the dimensions of the argument into the specified dimensions,
-- and just replicate the data along all other dimensions.
-- The list of dimensions indicies must have the same rank as the argument array
-- and it must be strictly ascending.
broadcast :: forall ds sh' sh v a .
             (Shape sh, Shape sh',
              Broadcast ds sh sh',
              Vector v, VecElem v a) =>
             Array sh v a -> Array sh' v a
broadcast :: Array sh v a -> Array sh' v a
broadcast Array sh v a
a = [Bool] -> Array Any v a -> Array sh' v a
forall (sh :: [Nat]) (v :: * -> *) a (sh' :: [Nat]).
[Bool] -> Array sh v a -> Array sh' v a
stretch' [Bool]
bc (Array Any v a -> Array sh' v a) -> Array Any v a -> Array sh' v a
forall a b. (a -> b) -> a -> b
$
              ShapeL -> ShapeL -> Array sh v a -> Array Any v a
forall (v :: * -> *) a (sh :: [Nat]) (sh' :: [Nat]).
(Vector v, VecElem v a) =>
ShapeL -> ShapeL -> Array sh v a -> Array sh' v a
reshape' ShapeL
sh ShapeL
rsh Array sh v a
a
  where sh' :: ShapeL
sh' = Proxy sh' -> ShapeL
forall (s :: [Nat]). Shape s => Proxy s -> ShapeL
shapeP (Proxy sh'
forall k (t :: k). Proxy t
Proxy :: Proxy sh')
        sh :: ShapeL
sh = Proxy sh -> ShapeL
forall (s :: [Nat]). Shape s => Proxy s -> ShapeL
shapeP (Proxy sh
forall k (t :: k). Proxy t
Proxy :: Proxy sh)
        rsh :: ShapeL
rsh = [ if Bool
b then Int
1 else Int
s | (Int
s, Bool
b) <- ShapeL -> [Bool] -> [(Int, Bool)]
forall a b. [a] -> [b] -> [(a, b)]
zip ShapeL
sh' [Bool]
bc ]
        bc :: [Bool]
bc = Broadcast ds sh sh' => [Bool]
forall (ds :: [Nat]) (sh :: [Nat]) (sh' :: [Nat]).
Broadcast ds sh sh' =>
[Bool]
broadcasting @ds @sh @sh'

-- | Generate an array with a function that computes the value for each index.
{-# INLINE generate #-}
generate :: forall sh v a .
            (Vector v, VecElem v a, Shape sh) =>
            ([Int] -> a) -> Array sh v a
generate :: (ShapeL -> a) -> Array sh v a
generate = T v a -> Array sh v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v a -> Array sh v a)
-> ((ShapeL -> a) -> T v a) -> (ShapeL -> a) -> Array sh v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShapeL -> (ShapeL -> a) -> T v a
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
ShapeL -> (ShapeL -> a) -> T v a
generateT (Proxy sh -> ShapeL
forall (s :: [Nat]). Shape s => Proxy s -> ShapeL
shapeP (Proxy sh
forall k (t :: k). Proxy t
Proxy :: Proxy sh))

-- | Iterate a function n times.
{-# INLINE iterateN #-}
iterateN :: forall n v a .
            (Vector v, VecElem v a, KnownNat n) =>
            (a -> a) -> a -> Array '[n] v a
iterateN :: (a -> a) -> a -> Array '[n] v a
iterateN a -> a
f = T v a -> Array '[n] v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v a -> Array '[n] v a) -> (a -> T v a) -> a -> Array '[n] v a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> (a -> a) -> a -> T v a
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
Int -> (a -> a) -> a -> T v a
iterateNT (forall i. (KnownNat n, Num i) => i
forall (n :: Nat) i. (KnownNat n, Num i) => i
valueOf @n) a -> a
f

-- | Generate a vector from 0 to n-1.
{-# INLINE iota #-}
iota :: forall n v a .
        (Vector v, VecElem v a, KnownNat n, Enum a, Num a) =>
        Array '[n] v a
iota :: Array '[n] v a
iota = T v a -> Array '[n] v a
forall (sh :: [Nat]) (v :: * -> *) a. T v a -> Array sh v a
A (T v a -> Array '[n] v a) -> T v a -> Array '[n] v a
forall a b. (a -> b) -> a -> b
$ Int -> T v a
forall (v :: * -> *) a.
(Vector v, VecElem v a, Enum a, Num a) =>
Int -> T v a
iotaT (forall i. (KnownNat n, Num i) => i
forall (n :: Nat) i. (KnownNat n, Num i) => i
valueOf @n)