-- 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.

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE ExplicitNamespaces #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.Array.Internal.ShapedS(
  Array(..), Shape(..), Size, Rank, Vector, ShapeL, Unbox,
  Window, Stride, Permute, Permutation, ValidDims,
  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 qualified Data.Vector.Storable as V
import GHC.Generics(Generic)
import GHC.TypeLits(KnownNat, type (+), type (<=))
import GHC.Stack(HasCallStack)
import Test.QuickCheck hiding (generate)
import Text.PrettyPrint.HughesPJClass hiding ((<>))

import Data.Array.Internal.DynamicS()  -- Vector instance
import qualified Data.Array.Internal.Shaped as S
import qualified Data.Array.Internal.ShapedG as G
import Data.Array.Internal(ShapeL, Vector)
import Data.Array.Internal.Shape

type Unbox = V.Storable

newtype Array sh a = A { Array sh a -> Array sh Vector a
unA :: G.Array sh V.Vector a }
  deriving (PrettyLevel -> [Array sh a] -> Doc
PrettyLevel -> Rational -> Array sh a -> Doc
Array sh a -> Doc
(PrettyLevel -> Rational -> Array sh a -> Doc)
-> (Array sh a -> Doc)
-> (PrettyLevel -> [Array sh a] -> Doc)
-> Pretty (Array sh a)
forall (sh :: [Nat]) a.
(Pretty a, Storable a, Shape sh) =>
PrettyLevel -> [Array sh a] -> Doc
forall (sh :: [Nat]) a.
(Pretty a, Storable a, Shape sh) =>
PrettyLevel -> Rational -> Array sh a -> Doc
forall (sh :: [Nat]) a.
(Pretty a, Storable a, Shape sh) =>
Array sh a -> Doc
forall a.
(PrettyLevel -> Rational -> a -> Doc)
-> (a -> Doc) -> (PrettyLevel -> [a] -> Doc) -> Pretty a
pPrintList :: PrettyLevel -> [Array sh a] -> Doc
$cpPrintList :: forall (sh :: [Nat]) a.
(Pretty a, Storable a, Shape sh) =>
PrettyLevel -> [Array sh a] -> Doc
pPrint :: Array sh a -> Doc
$cpPrint :: forall (sh :: [Nat]) a.
(Pretty a, Storable a, Shape sh) =>
Array sh a -> Doc
pPrintPrec :: PrettyLevel -> Rational -> Array sh a -> Doc
$cpPrintPrec :: forall (sh :: [Nat]) a.
(Pretty a, Storable a, Shape sh) =>
PrettyLevel -> Rational -> Array sh a -> Doc
Pretty, (forall x. Array sh a -> Rep (Array sh a) x)
-> (forall x. Rep (Array sh a) x -> Array sh a)
-> Generic (Array sh a)
forall (sh :: [Nat]) a x. Rep (Array sh a) x -> Array sh a
forall (sh :: [Nat]) a x. Array sh a -> Rep (Array sh a) x
forall x. Rep (Array sh a) x -> Array sh a
forall x. Array sh a -> Rep (Array sh a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall (sh :: [Nat]) a x. Rep (Array sh a) x -> Array sh a
$cfrom :: forall (sh :: [Nat]) a x. Array sh a -> Rep (Array sh a) x
Generic, Typeable (Array sh a)
DataType
Constr
Typeable (Array sh a)
-> (forall (c :: * -> *).
    (forall d b. Data d => c (d -> b) -> d -> c b)
    -> (forall g. g -> c g) -> Array sh a -> c (Array sh a))
-> (forall (c :: * -> *).
    (forall b r. Data b => c (b -> r) -> c r)
    -> (forall r. r -> c r) -> Constr -> c (Array sh a))
-> (Array sh a -> Constr)
-> (Array sh a -> DataType)
-> (forall (t :: * -> *) (c :: * -> *).
    Typeable t =>
    (forall d. Data d => c (t d)) -> Maybe (c (Array sh a)))
-> (forall (t :: * -> * -> *) (c :: * -> *).
    Typeable t =>
    (forall d e. (Data d, Data e) => c (t d e))
    -> Maybe (c (Array sh a)))
-> ((forall b. Data b => b -> b) -> Array sh a -> Array sh a)
-> (forall r r'.
    (r -> r' -> r)
    -> r -> (forall d. Data d => d -> r') -> Array sh a -> r)
-> (forall r r'.
    (r' -> r -> r)
    -> r -> (forall d. Data d => d -> r') -> Array sh a -> r)
-> (forall u. (forall d. Data d => d -> u) -> Array sh a -> [u])
-> (forall u.
    Int -> (forall d. Data d => d -> u) -> Array sh a -> u)
-> (forall (m :: * -> *).
    Monad m =>
    (forall d. Data d => d -> m d) -> Array sh a -> m (Array sh a))
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> Array sh a -> m (Array sh a))
-> (forall (m :: * -> *).
    MonadPlus m =>
    (forall d. Data d => d -> m d) -> Array sh a -> m (Array sh a))
-> Data (Array sh a)
Array sh a -> DataType
Array sh a -> Constr
(forall b. Data b => b -> b) -> Array sh a -> Array sh a
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Array sh a -> c (Array sh a)
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Array sh a)
forall (sh :: [Nat]) a.
(Typeable sh, Data a, Storable a) =>
Typeable (Array sh a)
forall (sh :: [Nat]) a.
(Typeable sh, Data a, Storable a) =>
Array sh a -> DataType
forall (sh :: [Nat]) a.
(Typeable sh, Data a, Storable a) =>
Array sh a -> Constr
forall (sh :: [Nat]) a.
(Typeable sh, Data a, Storable a) =>
(forall b. Data b => b -> b) -> Array sh a -> Array sh a
forall (sh :: [Nat]) a u.
(Typeable sh, Data a, Storable a) =>
Int -> (forall d. Data d => d -> u) -> Array sh a -> u
forall (sh :: [Nat]) a u.
(Typeable sh, Data a, Storable a) =>
(forall d. Data d => d -> u) -> Array sh a -> [u]
forall (sh :: [Nat]) a r r'.
(Typeable sh, Data a, Storable a) =>
(r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh a -> r
forall (sh :: [Nat]) a r r'.
(Typeable sh, Data a, Storable a) =>
(r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh a -> r
forall (sh :: [Nat]) a (m :: * -> *).
(Typeable sh, Data a, Storable a, Monad m) =>
(forall d. Data d => d -> m d) -> Array sh a -> m (Array sh a)
forall (sh :: [Nat]) a (m :: * -> *).
(Typeable sh, Data a, Storable a, MonadPlus m) =>
(forall d. Data d => d -> m d) -> Array sh a -> m (Array sh a)
forall (sh :: [Nat]) a (c :: * -> *).
(Typeable sh, Data a, Storable a) =>
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Array sh a)
forall (sh :: [Nat]) a (c :: * -> *).
(Typeable sh, Data a, Storable a) =>
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Array sh a -> c (Array sh a)
forall (sh :: [Nat]) a (t :: * -> *) (c :: * -> *).
(Typeable sh, Data a, Storable a, Typeable t) =>
(forall d. Data d => c (t d)) -> Maybe (c (Array sh a))
forall (sh :: [Nat]) a (t :: * -> * -> *) (c :: * -> *).
(Typeable sh, Data a, Storable a, Typeable t) =>
(forall d e. (Data d, Data e) => c (t d e))
-> Maybe (c (Array sh 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 a -> u
forall u. (forall d. Data d => d -> u) -> Array sh a -> [u]
forall r r'.
(r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh a -> r
forall r r'.
(r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh a -> r
forall (m :: * -> *).
Monad m =>
(forall d. Data d => d -> m d) -> Array sh a -> m (Array sh a)
forall (m :: * -> *).
MonadPlus m =>
(forall d. Data d => d -> m d) -> Array sh a -> m (Array sh a)
forall (c :: * -> *).
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Array sh a)
forall (c :: * -> *).
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Array sh a -> c (Array sh a)
forall (t :: * -> *) (c :: * -> *).
Typeable t =>
(forall d. Data d => c (t d)) -> Maybe (c (Array sh a))
forall (t :: * -> * -> *) (c :: * -> *).
Typeable t =>
(forall d e. (Data d, Data e) => c (t d e))
-> Maybe (c (Array sh a))
$cA :: Constr
$tArray :: DataType
gmapMo :: (forall d. Data d => d -> m d) -> Array sh a -> m (Array sh a)
$cgmapMo :: forall (sh :: [Nat]) a (m :: * -> *).
(Typeable sh, Data a, Storable a, MonadPlus m) =>
(forall d. Data d => d -> m d) -> Array sh a -> m (Array sh a)
gmapMp :: (forall d. Data d => d -> m d) -> Array sh a -> m (Array sh a)
$cgmapMp :: forall (sh :: [Nat]) a (m :: * -> *).
(Typeable sh, Data a, Storable a, MonadPlus m) =>
(forall d. Data d => d -> m d) -> Array sh a -> m (Array sh a)
gmapM :: (forall d. Data d => d -> m d) -> Array sh a -> m (Array sh a)
$cgmapM :: forall (sh :: [Nat]) a (m :: * -> *).
(Typeable sh, Data a, Storable a, Monad m) =>
(forall d. Data d => d -> m d) -> Array sh a -> m (Array sh a)
gmapQi :: Int -> (forall d. Data d => d -> u) -> Array sh a -> u
$cgmapQi :: forall (sh :: [Nat]) a u.
(Typeable sh, Data a, Storable a) =>
Int -> (forall d. Data d => d -> u) -> Array sh a -> u
gmapQ :: (forall d. Data d => d -> u) -> Array sh a -> [u]
$cgmapQ :: forall (sh :: [Nat]) a u.
(Typeable sh, Data a, Storable a) =>
(forall d. Data d => d -> u) -> Array sh a -> [u]
gmapQr :: (r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh a -> r
$cgmapQr :: forall (sh :: [Nat]) a r r'.
(Typeable sh, Data a, Storable a) =>
(r' -> r -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh a -> r
gmapQl :: (r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh a -> r
$cgmapQl :: forall (sh :: [Nat]) a r r'.
(Typeable sh, Data a, Storable a) =>
(r -> r' -> r)
-> r -> (forall d. Data d => d -> r') -> Array sh a -> r
gmapT :: (forall b. Data b => b -> b) -> Array sh a -> Array sh a
$cgmapT :: forall (sh :: [Nat]) a.
(Typeable sh, Data a, Storable a) =>
(forall b. Data b => b -> b) -> Array sh a -> Array sh a
dataCast2 :: (forall d e. (Data d, Data e) => c (t d e))
-> Maybe (c (Array sh a))
$cdataCast2 :: forall (sh :: [Nat]) a (t :: * -> * -> *) (c :: * -> *).
(Typeable sh, Data a, Storable a, Typeable t) =>
(forall d e. (Data d, Data e) => c (t d e))
-> Maybe (c (Array sh a))
dataCast1 :: (forall d. Data d => c (t d)) -> Maybe (c (Array sh a))
$cdataCast1 :: forall (sh :: [Nat]) a (t :: * -> *) (c :: * -> *).
(Typeable sh, Data a, Storable a, Typeable t) =>
(forall d. Data d => c (t d)) -> Maybe (c (Array sh a))
dataTypeOf :: Array sh a -> DataType
$cdataTypeOf :: forall (sh :: [Nat]) a.
(Typeable sh, Data a, Storable a) =>
Array sh a -> DataType
toConstr :: Array sh a -> Constr
$ctoConstr :: forall (sh :: [Nat]) a.
(Typeable sh, Data a, Storable a) =>
Array sh a -> Constr
gunfold :: (forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Array sh a)
$cgunfold :: forall (sh :: [Nat]) a (c :: * -> *).
(Typeable sh, Data a, Storable a) =>
(forall b r. Data b => c (b -> r) -> c r)
-> (forall r. r -> c r) -> Constr -> c (Array sh a)
gfoldl :: (forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Array sh a -> c (Array sh a)
$cgfoldl :: forall (sh :: [Nat]) a (c :: * -> *).
(Typeable sh, Data a, Storable a) =>
(forall d b. Data d => c (d -> b) -> d -> c b)
-> (forall g. g -> c g) -> Array sh a -> c (Array sh a)
$cp1Data :: forall (sh :: [Nat]) a.
(Typeable sh, Data a, Storable a) =>
Typeable (Array sh a)
Data)

instance NFData a => NFData (Array sh a)

instance (Unbox a, Show a, Shape sh) => Show (Array sh a) where
  showsPrec :: Int -> Array sh a -> ShowS
showsPrec Int
p = Int -> Array sh Vector a -> ShowS
forall a. Show a => Int -> a -> ShowS
showsPrec Int
p (Array sh Vector a -> ShowS)
-> (Array sh a -> Array sh Vector a) -> Array sh a -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

instance (Read a, Unbox a, Shape sh) => Read (Array sh a) where
  readsPrec :: Int -> ReadS (Array sh a)
readsPrec Int
p String
s = [(Array sh Vector a -> Array sh a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A Array sh Vector a
a, String
r) | (Array sh Vector a
a, String
r) <- Int -> ReadS (Array sh Vector a)
forall a. Read a => Int -> ReadS a
readsPrec Int
p String
s]

instance Eq (G.Array sh V.Vector a) => Eq (Array sh a) where
  Array sh a
x == :: Array sh a -> Array sh a -> Bool
== Array sh a
y = Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA Array sh a
x Array sh Vector a -> Array sh Vector a -> Bool
forall a. Eq a => a -> a -> Bool
== Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA Array sh a
y
  {-# INLINE (==) #-}

instance Ord (G.Array sh V.Vector a) => Ord (Array sh a) where
  compare :: Array sh a -> Array sh a -> Ordering
compare Array sh a
x Array sh a
y = Array sh Vector a -> Array sh Vector a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA Array sh a
x) (Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA Array sh a
y)
  {-# INLINE compare #-}

-- | The number of elements in the array.
{-# INLINE size #-}
size :: forall sh a . (Shape sh) => Array sh a -> Int
size :: Array sh a -> Int
size = Array sh Vector a -> Int
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> Int
G.size (Array sh Vector a -> Int)
-> (Array sh a -> Array sh Vector a) -> Array sh a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | 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.
shapeL :: (Shape sh) => Array sh a -> ShapeL
shapeL :: Array sh a -> ShapeL
shapeL = Array sh Vector a -> ShapeL
forall (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array sh v a -> ShapeL
G.shapeL (Array sh Vector a -> ShapeL)
-> (Array sh a -> Array sh Vector a) -> Array sh a -> ShapeL
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | The rank of an array, i.e., the number if dimensions it has,
-- which is the @n@ in @Array n a@.
-- O(1) time.
rank :: (Shape sh, KnownNat (Rank sh)) => Array sh a -> Int
rank :: Array sh a -> Int
rank = Array sh Vector a -> Int
forall (sh :: [Nat]) (v :: * -> *) a.
(Shape sh, KnownNat (Rank sh)) =>
Array sh v a -> Int
G.rank (Array sh Vector a -> Int)
-> (Array sh a -> Array sh Vector a) -> Array sh a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Index into an array.  Fails if the index is out of bounds.
-- O(1) time.
index :: (Unbox a, KnownNat s) => Array (s:sh) a -> Int -> Array sh a
index :: Array (s : sh) a -> Int -> Array sh a
index Array (s : sh) a
a = Array sh Vector a -> Array sh a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh Vector a -> Array sh a)
-> (Int -> Array sh Vector a) -> Int -> Array sh a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array (s : sh) Vector a -> Int -> Array sh Vector a
forall (s :: Nat) (sh :: [Nat]) (v :: * -> *) a.
(HasCallStack, Vector v, KnownNat s) =>
Array (s : sh) v a -> Int -> Array sh v a
G.index (Array (s : sh) a -> Array (s : sh) Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA Array (s : sh) a
a)

-- | Convert to a list with the elements in the linearization order.
-- O(n) time.
toList :: forall sh a . (Unbox a, Shape sh) => Array sh a -> [a]
toList :: Array sh a -> [a]
toList = Array sh Vector a -> [a]
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Shape sh) =>
Array sh v a -> [a]
G.toList (Array sh Vector a -> [a])
-> (Array sh a -> Array sh Vector a) -> Array sh a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | 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.
fromList :: forall sh a . (HasCallStack, Unbox a, Shape sh) => [a] -> Array sh a
fromList :: [a] -> Array sh a
fromList = Array sh Vector a -> Array sh a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh Vector a -> Array sh a)
-> ([a] -> Array sh Vector a) -> [a] -> Array sh a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Array sh Vector a
forall (sh :: [Nat]) (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a, Shape sh) =>
[a] -> Array sh v a
G.fromList

-- | 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).
toVector :: (Unbox a, Shape sh) => Array sh a -> V.Vector a
toVector :: Array sh a -> Vector a
toVector = Array sh Vector a -> Vector a
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Shape sh) =>
Array sh v a -> v a
G.toVector (Array sh Vector a -> Vector a)
-> (Array sh a -> Array sh Vector a) -> Array sh a -> Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | 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.
fromVector :: forall sh a . (HasCallStack, Unbox a, Shape sh) => V.Vector a -> Array sh a
fromVector :: Vector a -> Array sh a
fromVector = Array sh Vector a -> Array sh a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh Vector a -> Array sh a)
-> (Vector a -> Array sh Vector a) -> Vector a -> Array sh a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector a -> Array sh Vector a
forall (sh :: [Nat]) (v :: * -> *) a.
(HasCallStack, Vector v, VecElem v a, Shape sh) =>
v a -> Array sh v a
G.fromVector

-- | 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.
normalize :: (Unbox a, Shape sh) => Array sh a -> Array sh a
normalize :: Array sh a -> Array sh a
normalize = Array sh Vector a -> Array sh a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh Vector a -> Array sh a)
-> (Array sh a -> Array sh Vector a) -> Array sh a -> Array sh a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh Vector a -> Array sh Vector a
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Shape sh) =>
Array sh v a -> Array sh v a
G.normalize (Array sh Vector a -> Array sh Vector a)
-> (Array sh a -> Array sh Vector a)
-> Array sh a
-> Array sh Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Change the shape of an array.  Fails if the arrays have different number of elements.
-- O(n) or O(1) time.
reshape :: forall sh' sh a . (Unbox a, Shape sh, Shape sh', Size sh ~ Size sh') =>
           Array sh a -> Array sh' a
reshape :: Array sh a -> Array sh' a
reshape = Array sh' Vector a -> Array sh' a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh' Vector a -> Array sh' a)
-> (Array sh a -> Array sh' Vector a) -> Array sh a -> Array sh' a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh Vector a -> Array sh' Vector a
forall (sh' :: [Nat]) (sh :: [Nat]) (v :: * -> *) a.
(Vector v, VecElem v a, Shape sh, Shape sh', Size sh ~ Size sh') =>
Array sh v a -> Array sh' v a
G.reshape (Array sh Vector a -> Array sh' Vector a)
-> (Array sh a -> Array sh Vector a)
-> Array sh a
-> Array sh' Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | 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.
stretch :: forall sh' sh a . (Shape sh, Shape sh', ValidStretch sh sh') => Array sh a -> Array sh' a
stretch :: Array sh a -> Array sh' a
stretch = Array sh' Vector a -> Array sh' a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh' Vector a -> Array sh' a)
-> (Array sh a -> Array sh' Vector a) -> Array sh a -> Array sh' a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh Vector a -> Array sh' Vector a
forall (sh' :: [Nat]) (sh :: [Nat]) (v :: * -> *) a.
(Shape sh, Shape sh', ValidStretch sh sh') =>
Array sh v a -> Array sh' v a
G.stretch (Array sh Vector a -> Array sh' Vector a)
-> (Array sh a -> Array sh Vector a)
-> Array sh a
-> Array sh' Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Change the size of the outermost dimension by replication.
stretchOuter :: (KnownNat s, Shape sh) => Array (1 : sh) a -> Array (s : sh) a
stretchOuter :: Array (1 : sh) a -> Array (s : sh) a
stretchOuter = Array (s : sh) Vector a -> Array (s : sh) a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array (s : sh) Vector a -> Array (s : sh) a)
-> (Array (1 : sh) a -> Array (s : sh) Vector a)
-> Array (1 : sh) a
-> Array (s : sh) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array (1 : sh) Vector a -> Array (s : sh) Vector a
forall (s :: Nat) (sh :: [Nat]) (v :: * -> *) a.
Shape sh =>
Array (1 : sh) v a -> Array (s : sh) v a
G.stretchOuter (Array (1 : sh) Vector a -> Array (s : sh) Vector a)
-> (Array (1 : sh) a -> Array (1 : sh) Vector a)
-> Array (1 : sh) a
-> Array (s : sh) Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array (1 : sh) a -> Array (1 : sh) Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

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

-- | Convert a scalar (rank 0) array to a value.
-- O(1) time.
unScalar :: (Unbox a) => Array '[] a -> a
unScalar :: Array '[] a -> a
unScalar = Array '[] Vector a -> a
forall (v :: * -> *) a.
(Vector v, VecElem v a) =>
Array '[] v a -> a
G.unScalar (Array '[] Vector a -> a)
-> (Array '[] a -> Array '[] Vector a) -> Array '[] a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array '[] a -> Array '[] Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Make an array with all elements having the same value.
-- O(1) time.
constant :: forall sh a . (Unbox a, Shape sh) =>
            a -> Array sh a
constant :: a -> Array sh a
constant = Array sh Vector a -> Array sh a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh Vector a -> Array sh a)
-> (a -> Array sh Vector a) -> a -> Array sh a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Array sh Vector a
forall (sh :: [Nat]) (v :: * -> *) a.
(Vector v, VecElem v a, Shape sh) =>
a -> Array sh v a
G.constant

-- | Map over the array elements.
-- O(n) time.
mapA :: (Unbox a, Unbox b, Shape sh) => (a -> b) -> Array sh a -> Array sh b
mapA :: (a -> b) -> Array sh a -> Array sh b
mapA a -> b
f = Array sh Vector b -> Array sh b
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh Vector b -> Array sh b)
-> (Array sh a -> Array sh Vector b) -> Array sh a -> Array sh b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> Array sh Vector a -> Array sh Vector b
forall (v :: * -> *) a b (sh :: [Nat]).
(Vector v, VecElem v a, VecElem v b, Shape sh) =>
(a -> b) -> Array sh v a -> Array sh v b
G.mapA a -> b
f (Array sh Vector a -> Array sh Vector b)
-> (Array sh a -> Array sh Vector a)
-> Array sh a
-> Array sh Vector b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Map over the array elements.
-- O(n) time.
zipWithA :: (Unbox a, Unbox b, Unbox c, Shape sh) => (a -> b -> c) -> Array sh a -> Array sh b -> Array sh c
zipWithA :: (a -> b -> c) -> Array sh a -> Array sh b -> Array sh c
zipWithA a -> b -> c
f Array sh a
a Array sh b
b = Array sh Vector c -> Array sh c
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh Vector c -> Array sh c)
-> Array sh Vector c -> Array sh c
forall a b. (a -> b) -> a -> b
$ (a -> b -> c)
-> Array sh Vector a -> Array sh Vector b -> Array sh Vector c
forall (v :: * -> *) a b c (sh :: [Nat]).
(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
G.zipWithA a -> b -> c
f (Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA Array sh a
a) (Array sh b -> Array sh Vector b
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA Array sh b
b)

-- | Map over the array elements.
-- O(n) time.
zipWith3A :: (Unbox a, Unbox b, Unbox c, Unbox d, Shape sh) => (a -> b -> c -> d) -> Array sh a -> Array sh b -> Array sh c -> Array sh d
zipWith3A :: (a -> b -> c -> d)
-> Array sh a -> Array sh b -> Array sh c -> Array sh d
zipWith3A a -> b -> c -> d
f Array sh a
a Array sh b
b Array sh c
c = Array sh Vector d -> Array sh d
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh Vector d -> Array sh d)
-> Array sh Vector d -> Array sh d
forall a b. (a -> b) -> a -> b
$ (a -> b -> c -> d)
-> Array sh Vector a
-> Array sh Vector b
-> Array sh Vector c
-> Array sh Vector d
forall (v :: * -> *) a b c d (sh :: [Nat]).
(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
G.zipWith3A a -> b -> c -> d
f (Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA Array sh a
a) (Array sh b -> Array sh Vector b
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA Array sh b
b) (Array sh c -> Array sh Vector c
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA Array sh c
c)

-- | Pad each dimension on the low and high side with the given value.
-- O(n) time.
pad :: forall ps sh' sh a . (Unbox a, Padded ps sh sh', Shape sh) =>
       a -> Array sh a -> Array sh' a
pad :: a -> Array sh a -> Array sh' a
pad a
v = Array sh' Vector a -> Array sh' a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh' Vector a -> Array sh' a)
-> (Array sh a -> Array sh' Vector a) -> Array sh a -> Array sh' a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Array sh Vector a -> Array sh' Vector a
forall (ps :: [(Nat, Nat)]) (sh' :: [Nat]) (sh :: [Nat]) a
       (v :: * -> *).
(HasCallStack, Vector v, VecElem v a, Padded ps sh sh',
 Shape sh) =>
a -> Array sh v a -> Array sh' v a
G.pad @ps a
v (Array sh Vector a -> Array sh' Vector a)
-> (Array sh a -> Array sh Vector a)
-> Array sh a
-> Array sh' Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | 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.
transpose :: forall is sh a .
             (Permutation is, Rank is <= Rank sh, Shape sh, Shape is, KnownNat (Rank sh)) =>
             Array sh a -> Array (Permute is sh) a
transpose :: Array sh a -> Array (Permute is sh) a
transpose = Array (Permute is sh) Vector a -> Array (Permute is sh) a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array (Permute is sh) Vector a -> Array (Permute is sh) a)
-> (Array sh a -> Array (Permute is sh) Vector a)
-> Array sh a
-> Array (Permute is sh) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (is :: [Nat]) (sh :: [Nat]) (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
forall (sh :: [Nat]) (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
G.transpose @is (Array sh Vector a -> Array (Permute is sh) Vector a)
-> (Array sh a -> Array sh Vector a)
-> Array sh a
-> Array (Permute is sh) Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Append two arrays along the outermost dimension.
-- All dimensions, except the outermost, must be the same.
-- O(n) time.
append :: (Unbox a, Shape sh, KnownNat m, KnownNat n, KnownNat (m+n)) =>
          Array (m ': sh) a -> Array (n ': sh) a -> Array (m+n ': sh) a
append :: Array (m : sh) a -> Array (n : sh) a -> Array ((m + n) : sh) a
append Array (m : sh) a
x Array (n : sh) a
y = Array ((m + n) : sh) Vector a -> Array ((m + n) : sh) a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array ((m + n) : sh) Vector a -> Array ((m + n) : sh) a)
-> Array ((m + n) : sh) Vector a -> Array ((m + n) : sh) a
forall a b. (a -> b) -> a -> b
$ Array (m : sh) Vector a
-> Array (n : sh) Vector a -> Array ((m + n) : sh) Vector a
forall (v :: * -> *) a (sh :: [Nat]) (m :: Nat) (n :: Nat).
(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
G.append (Array (m : sh) a -> Array (m : sh) Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA Array (m : sh) a
x) (Array (n : sh) a -> Array (n : sh) Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA Array (n : sh) a
y)

-- | 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.
ravel :: (Unbox a, Shape sh, KnownNat s) =>
         S.Array '[s] (Array sh a) -> Array (s:sh) a
ravel :: Array '[s] (Array sh a) -> Array (s : sh) a
ravel = Array (s : sh) Vector a -> Array (s : sh) a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array (s : sh) Vector a -> Array (s : sh) a)
-> (Array '[s] (Array sh a) -> Array (s : sh) Vector a)
-> Array '[s] (Array sh a)
-> Array (s : sh) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array '[s] Vector (Array sh Vector a) -> Array (s : sh) Vector a
forall (v :: * -> *) (v' :: * -> *) a (sh :: [Nat]) (s :: Nat).
(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
G.ravel (Array '[s] Vector (Array sh Vector a) -> Array (s : sh) Vector a)
-> (Array '[s] (Array sh a)
    -> Array '[s] Vector (Array sh Vector a))
-> Array '[s] (Array sh a)
-> Array (s : sh) Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array sh a -> Array sh Vector a)
-> Array '[s] Vector (Array sh a)
-> Array '[s] Vector (Array sh Vector a)
forall (v :: * -> *) a b (sh :: [Nat]).
(Vector v, VecElem v a, VecElem v b, Shape sh) =>
(a -> b) -> Array sh v a -> Array sh v b
G.mapA Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA (Array '[s] Vector (Array sh a)
 -> Array '[s] Vector (Array sh Vector a))
-> (Array '[s] (Array sh a) -> Array '[s] Vector (Array sh a))
-> Array '[s] (Array sh a)
-> Array '[s] Vector (Array sh Vector a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array '[s] (Array sh a) -> Array '[s] Vector (Array sh a)
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
S.unA

-- | Turn an array into a nested array, this is the inverse of 'ravel'.
-- I.e., @ravel . unravel == id@.
unravel :: (Unbox a, Shape sh, KnownNat s) =>
           Array (s:sh) a -> S.Array '[s] (Array sh a)
unravel :: Array (s : sh) a -> Array '[s] (Array sh a)
unravel = Array '[s] Vector (Array sh a) -> Array '[s] (Array sh a)
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
S.A (Array '[s] Vector (Array sh a) -> Array '[s] (Array sh a))
-> (Array (s : sh) a -> Array '[s] Vector (Array sh a))
-> Array (s : sh) a
-> Array '[s] (Array sh a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array sh Vector a -> Array sh a)
-> Array '[s] Vector (Array sh Vector a)
-> Array '[s] Vector (Array sh a)
forall (v :: * -> *) a b (sh :: [Nat]).
(Vector v, VecElem v a, VecElem v b, Shape sh) =>
(a -> b) -> Array sh v a -> Array sh v b
G.mapA Array sh Vector a -> Array sh a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array '[s] Vector (Array sh Vector a)
 -> Array '[s] Vector (Array sh a))
-> (Array (s : sh) a -> Array '[s] Vector (Array sh Vector a))
-> Array (s : sh) a
-> Array '[s] Vector (Array sh a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array (s : sh) Vector a -> Array '[s] Vector (Array sh Vector a)
forall (v :: * -> *) (v' :: * -> *) a (sh :: [Nat]) (s :: Nat).
(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)
G.unravel (Array (s : sh) Vector a -> Array '[s] Vector (Array sh Vector a))
-> (Array (s : sh) a -> Array (s : sh) Vector a)
-> Array (s : sh) a
-> Array '[s] Vector (Array sh Vector a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array (s : sh) a -> Array (s : sh) Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | 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]@.
-- O(1) time.
window :: forall ws sh' sh a .
          (Window ws sh sh', KnownNat (Rank ws)) =>
          Array sh a -> Array sh' a
window :: Array sh a -> Array sh' a
window = Array sh' Vector a -> Array sh' a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh' Vector a -> Array sh' a)
-> (Array sh a -> Array sh' Vector a) -> Array sh a -> Array sh' a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (ws :: [Nat]) (sh' :: [Nat]) (sh :: [Nat]) (v :: * -> *) a.
(Window ws sh sh', Vector v, KnownNat (Rank ws)) =>
Array sh v a -> Array sh' v a
forall (sh' :: [Nat]) (sh :: [Nat]) (v :: * -> *) a.
(Window ws sh sh', Vector v, KnownNat (Rank ws)) =>
Array sh v a -> Array sh' v a
G.window @ws (Array sh Vector a -> Array sh' Vector a)
-> (Array sh a -> Array sh Vector a)
-> Array sh a
-> Array sh' Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | 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.
stride :: forall ts sh' sh a .
          (Stride ts sh sh', Shape ts) =>
          Array sh a -> Array sh' a
stride :: Array sh a -> Array sh' a
stride = Array sh' Vector a -> Array sh' a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh' Vector a -> Array sh' a)
-> (Array sh a -> Array sh' Vector a) -> Array sh a -> Array sh' a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (ts :: [Nat]) (sh' :: [Nat]) (sh :: [Nat]) (v :: * -> *) a.
(Stride ts sh sh', Vector v, Shape ts) =>
Array sh v a -> Array sh' v a
forall (sh' :: [Nat]) (sh :: [Nat]) (v :: * -> *) a.
(Stride ts sh sh', Vector v, Shape ts) =>
Array sh v a -> Array sh' v a
G.stride @ts (Array sh Vector a -> Array sh' Vector a)
-> (Array sh a -> Array sh Vector a)
-> Array sh a
-> Array sh' Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Extract a slice of an array.
-- The first argument is a list of (offset, length) pairs.
-- The length of the slicing argument must not exceed the rank of the arrar.
-- The extracted slice mul fall within the array dimensions.
-- E.g. @slice [1,2] (fromList [4] [1,2,3,4]) == [2,3]@.
-- O(1) time.
slice :: forall sl sh' sh a .
         (Slice sl sh sh') =>
         Array sh a -> Array sh' a
slice :: Array sh a -> Array sh' a
slice = Array sh' Vector a -> Array sh' a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh' Vector a -> Array sh' a)
-> (Array sh a -> Array sh' Vector a) -> Array sh a -> Array sh' a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (sh' :: [Nat]) (sh :: [Nat]) (v :: * -> *) a.
Slice sl sh sh' =>
Array sh v a -> Array sh' v a
forall (sl :: [(Nat, Nat)]) (sh' :: [Nat]) (sh :: [Nat])
       (v :: * -> *) a.
Slice sl sh sh' =>
Array sh v a -> Array sh' v a
G.slice @sl (Array sh Vector a -> Array sh' Vector a)
-> (Array sh a -> Array sh Vector a)
-> Array sh a
-> Array sh' Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | 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.
rerank :: forall n i o sh a b .
          (Unbox a, Unbox b,
           Drop n sh ~ i, Shape sh, KnownNat n, Shape o, Shape (Take n sh ++ o)) =>
          (Array i a -> Array o b) -> Array sh a -> Array (Take n sh ++ o) b
rerank :: (Array i a -> Array o b) -> Array sh a -> Array (Take n sh ++ o) b
rerank Array i a -> Array o b
f = Array (Take n sh ++ o) Vector b -> Array (Take n sh ++ o) b
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array (Take n sh ++ o) Vector b -> Array (Take n sh ++ o) b)
-> (Array sh a -> Array (Take n sh ++ o) Vector b)
-> Array sh a
-> Array (Take n sh ++ o) b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Array i Vector a -> Array o Vector b)
-> Array sh Vector a -> Array (Take n sh ++ o) Vector b
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
G.rerank @n (Array o b -> Array o Vector b
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA (Array o b -> Array o Vector b)
-> (Array i Vector a -> Array o b)
-> Array i Vector a
-> Array o Vector b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array i a -> Array o b
f (Array i a -> Array o b)
-> (Array i Vector a -> Array i a) -> Array i Vector a -> Array o b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array i Vector a -> Array i a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A) (Array sh Vector a -> Array (Take n sh ++ o) Vector b)
-> (Array sh a -> Array sh Vector a)
-> Array sh a
-> Array (Take n sh ++ o) Vector b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | 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.
rerank2 :: forall n i o sh a b c .
           (Unbox a, Unbox b, Unbox c,
            Drop n sh ~ i, Shape sh, KnownNat n, Shape o, Shape (Take n sh ++ o)) =>
           (Array i a -> Array i b -> Array o c) -> Array sh a -> Array sh b -> Array (Take n sh ++ o) c
rerank2 :: (Array i a -> Array i b -> Array o c)
-> Array sh a -> Array sh b -> Array (Take n sh ++ o) c
rerank2 Array i a -> Array i b -> Array o c
f Array sh a
ta Array sh b
tb = Array (Take n sh ++ o) Vector c -> Array (Take n sh ++ o) c
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array (Take n sh ++ o) Vector c -> Array (Take n sh ++ o) c)
-> Array (Take n sh ++ o) Vector c -> Array (Take n sh ++ o) c
forall a b. (a -> b) -> a -> b
$ (Array i Vector a -> Array i Vector b -> Array o Vector c)
-> Array sh Vector a
-> Array sh Vector b
-> Array (Take n sh ++ o) Vector c
forall (n :: Nat) (i1 :: [Nat]) (i2 :: [Nat]) (o :: [Nat])
       (sh1 :: [Nat]) (sh2 :: [Nat]) (r :: [Nat]) (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
G.rerank2 @n (\ Array i Vector a
a Array i Vector b
b -> Array o c -> Array o Vector c
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA (Array o c -> Array o Vector c) -> Array o c -> Array o Vector c
forall a b. (a -> b) -> a -> b
$ Array i a -> Array i b -> Array o c
f (Array i Vector a -> Array i a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A Array i Vector a
a) (Array i Vector b -> Array i b
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A Array i Vector b
b)) (Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA Array sh a
ta) (Array sh b -> Array sh Vector b
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA Array sh b
tb)

-- | Reverse the given dimensions, with the outermost being dimension 0.
-- O(1) time.
rev :: forall rs sh a . (ValidDims rs sh, Shape rs, Shape sh) =>
       Array sh a -> Array sh a
rev :: Array sh a -> Array sh a
rev = Array sh Vector a -> Array sh a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh Vector a -> Array sh a)
-> (Array sh a -> Array sh Vector a) -> Array sh a -> Array sh a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (rs :: [Nat]) (sh :: [Nat]) (v :: * -> *) a.
(ValidDims rs sh, Shape rs, Shape sh) =>
Array sh v a -> Array sh v a
forall (sh :: [Nat]) (v :: * -> *) a.
(ValidDims rs sh, Shape rs, Shape sh) =>
Array sh v a -> Array sh v a
G.rev @rs (Array sh Vector a -> Array sh Vector a)
-> (Array sh a -> Array sh Vector a)
-> Array sh a
-> Array sh Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Reduce all elements of an array into a rank 0 array.
-- To reduce parts use 'rerank' and 'transpose' together with 'reduce'.
-- O(n) time.
reduce :: (Unbox a, Shape sh) => (a -> a -> a) -> a -> Array sh a -> Array '[] a
reduce :: (a -> a -> a) -> a -> Array sh a -> Array '[] a
reduce a -> a -> a
f a
z = Array '[] Vector a -> Array '[] a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array '[] Vector a -> Array '[] a)
-> (Array sh a -> Array '[] Vector a) -> Array sh a -> Array '[] a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> a) -> a -> Array sh Vector a -> Array '[] Vector a
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Shape sh) =>
(a -> a -> a) -> a -> Array sh v a -> Array '[] v a
G.reduce a -> a -> a
f a
z (Array sh Vector a -> Array '[] Vector a)
-> (Array sh a -> Array sh Vector a)
-> Array sh a
-> Array '[] Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Constrained version of 'foldr' for Arrays.
foldrA :: (Unbox a, Shape sh) => (a -> b -> b) -> b -> Array sh a -> b
foldrA :: (a -> b -> b) -> b -> Array sh a -> b
foldrA a -> b -> b
f b
z = (a -> b -> b) -> b -> Array sh Vector a -> b
forall (v :: * -> *) a (sh :: [Nat]) b.
(Vector v, VecElem v a, Shape sh) =>
(a -> b -> b) -> b -> Array sh v a -> b
G.foldrA a -> b -> b
f b
z (Array sh Vector a -> b)
-> (Array sh a -> Array sh Vector a) -> Array sh a -> b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Constrained version of 'traverse' for Arrays.
traverseA
  :: (Unbox a, Unbox b, Applicative f, Shape sh)
  => (a -> f b) -> Array sh a -> f (Array sh b)
traverseA :: (a -> f b) -> Array sh a -> f (Array sh b)
traverseA a -> f b
f = (Array sh Vector b -> Array sh b)
-> f (Array sh Vector b) -> f (Array sh b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Array sh Vector b -> Array sh b
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (f (Array sh Vector b) -> f (Array sh b))
-> (Array sh a -> f (Array sh Vector b))
-> Array sh a
-> f (Array sh b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> f b) -> Array sh Vector a -> f (Array sh Vector b)
forall (v :: * -> *) a b (f :: * -> *) (sh :: [Nat]).
(Vector v, VecElem v a, VecElem v b, Applicative f, Shape sh) =>
(a -> f b) -> Array sh v a -> f (Array sh v b)
G.traverseA a -> f b
f (Array sh Vector a -> f (Array sh Vector b))
-> (Array sh a -> Array sh Vector a)
-> Array sh a
-> f (Array sh Vector b)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Check if all elements of the array are equal.
allSameA :: (Shape sh, Unbox a, Eq a) => Array sh a -> Bool
allSameA :: Array sh a -> Bool
allSameA = Array sh Vector a -> Bool
forall (sh :: [Nat]) (v :: * -> *) a.
(Shape sh, Vector v, VecElem v a, Eq a) =>
Array sh v a -> Bool
G.allSameA (Array sh Vector a -> Bool)
-> (Array sh a -> Array sh Vector a) -> Array sh a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

instance (Shape sh, Arbitrary a, Unbox a) => Arbitrary (Array sh a) where arbitrary :: Gen (Array sh a)
arbitrary = Array sh Vector a -> Array sh a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh Vector a -> Array sh a)
-> Gen (Array sh Vector a) -> Gen (Array sh a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Gen (Array sh Vector a)
forall a. Arbitrary a => Gen a
arbitrary

{-# INLINE sumA #-}
sumA :: (Unbox a, Num a, Shape sh) => Array sh a -> a
sumA :: Array sh a -> a
sumA = Array sh Vector a -> a
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Num a, Shape sh) =>
Array sh v a -> a
G.sumA (Array sh Vector a -> a)
-> (Array sh a -> Array sh Vector a) -> Array sh a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Product of all elements.
{-# INLINE productA #-}
productA :: (Unbox a, Num a, Shape sh) => Array sh a -> a
productA :: Array sh a -> a
productA = Array sh Vector a -> a
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Num a, Shape sh) =>
Array sh v a -> a
G.productA (Array sh Vector a -> a)
-> (Array sh a -> Array sh Vector a) -> Array sh a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Maximum of all elements.
{-# INLINE maximumA #-}
maximumA :: (Unbox a, Ord a, Shape sh, 1 <= Size sh) => Array sh a -> a
maximumA :: Array sh a -> a
maximumA = Array sh Vector a -> a
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Ord a, Shape sh, 1 <= Size sh) =>
Array sh v a -> a
G.maximumA (Array sh Vector a -> a)
-> (Array sh a -> Array sh Vector a) -> Array sh a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Minimum of all elements.
{-# INLINE minimumA #-}
minimumA :: (Unbox a, Ord a, Shape sh, 1 <= Size sh) => Array sh a -> a
minimumA :: Array sh a -> a
minimumA = Array sh Vector a -> a
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Ord a, Shape sh, 1 <= Size sh) =>
Array sh v a -> a
G.minimumA (Array sh Vector a -> a)
-> (Array sh a -> Array sh Vector a) -> Array sh a -> a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Test if the predicate holds for any element.
{-# INLINE anyA #-}
anyA :: (Shape sh, Unbox a) => (a -> Bool) -> Array sh a -> Bool
anyA :: (a -> Bool) -> Array sh a -> Bool
anyA a -> Bool
p = (a -> Bool) -> Array sh Vector a -> Bool
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Shape sh) =>
(a -> Bool) -> Array sh v a -> Bool
G.anyA a -> Bool
p (Array sh Vector a -> Bool)
-> (Array sh a -> Array sh Vector a) -> Array sh a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | Test if the predicate holds for all elements.
{-# INLINE allA #-}
allA :: (Shape sh, Unbox a) => (a -> Bool) -> Array sh a -> Bool
allA :: (a -> Bool) -> Array sh a -> Bool
allA a -> Bool
p = (a -> Bool) -> Array sh Vector a -> Bool
forall (v :: * -> *) a (sh :: [Nat]).
(Vector v, VecElem v a, Shape sh) =>
(a -> Bool) -> Array sh v a -> Bool
G.allA a -> Bool
p (Array sh Vector a -> Bool)
-> (Array sh a -> Array sh Vector a) -> Array sh a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

-- | 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.
{-# INLINE broadcast #-}
broadcast :: forall ds sh' sh a .
             (Unbox a, Shape sh, Shape sh',
              G.Broadcast ds sh sh') =>
             Array sh a -> Array sh' a
broadcast :: Array sh a -> Array sh' a
broadcast = Array sh' Vector a -> Array sh' a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array sh' Vector a -> Array sh' a)
-> (Array sh a -> Array sh' Vector a) -> Array sh a -> Array sh' a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (ds :: [Nat]) (sh' :: [Nat]) (sh :: [Nat]) (v :: * -> *) a.
(Shape sh, Shape sh', Broadcast ds sh sh', Vector v,
 VecElem v a) =>
Array sh v a -> Array sh' v a
forall (v :: * -> *) a.
(Shape sh, Shape sh', Broadcast ds sh sh', Vector v,
 VecElem v a) =>
Array sh v a -> Array sh' v a
G.broadcast @ds @sh' @sh (Array sh Vector a -> Array sh' Vector a)
-> (Array sh a -> Array sh Vector a)
-> Array sh a
-> Array sh' Vector a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array sh a -> Array sh Vector a
forall (sh :: [Nat]) a. Array sh a -> Array sh Vector a
unA

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

-- | Iterate a function n times.
{-# INLINE iterateN #-}
iterateN :: forall n a .
            (Unbox a, KnownNat n) => (a -> a) -> a -> Array '[n] a
iterateN :: (a -> a) -> a -> Array '[n] a
iterateN a -> a
f = Array '[n] Vector a -> Array '[n] a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A (Array '[n] Vector a -> Array '[n] a)
-> (a -> Array '[n] Vector a) -> a -> Array '[n] a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a) -> a -> Array '[n] Vector a
forall (n :: Nat) (v :: * -> *) a.
(Vector v, VecElem v a, KnownNat n) =>
(a -> a) -> a -> Array '[n] v a
G.iterateN a -> a
f

-- | Generate a vector from 0 to n-1.
{-# INLINE iota #-}
iota :: (KnownNat n, Unbox a, Enum a, Num a) => Array '[n] a
iota :: Array '[n] a
iota = Array '[n] Vector a -> Array '[n] a
forall (sh :: [Nat]) a. Array sh Vector a -> Array sh a
A Array '[n] Vector a
forall (n :: Nat) (v :: * -> *) a.
(Vector v, VecElem v a, KnownNat n, Enum a, Num a) =>
Array '[n] v a
G.iota