from tensor import Tensor, TensorShape, TensorSpec
from math import trunc, mod
from sys.info import simdwidthof
from algorithm import vectorize, parallelize
from random import rand, seed
from python import Python
struct myTensor[dtype: DType]:
var t: Tensor[dtype]
@always_inline
fn __init__(inout self, *dims: Int):
self.t = rand[dtype](TensorSpec(dtype, dims))
@always_inline
fn __init__(inout self, owned t: Tensor[dtype]):
self.t = t
fn mean(self) -> Self:
var new_tensor = Tensor[dtype](self.t.dim(0),1)
alias simd_width: Int = simdwidthof[dtype]()
@parameter
fn parallel_reduce_rows(idx1: Int)->None:
@parameter
fn vectorize_reduce_row[simd_width: Int](idx2: Int) -> None:
new_tensor[idx1] += self.t.simd_load[simd_width](idx1*self.t.dim(1)+idx2).reduce_add()
vectorize[2*simd_width,vectorize_reduce_row](self.t.dim(1))
new_tensor[idx1] /= self.t.dim(1)
parallelize[parallel_reduce_rows](self.t.dim(0),8)
return Self(new_tensor)
fn print(self, prec: Int=4)->None:
let t = self.t
let rank = t.rank()
if rank == 0:
print("Error: Nothing to print. Tensor rank = 0")
return
var dim0:Int=0
var dim1:Int=0
var dim2:Int=0
if rank==0 or rank>3:
print("Error: Tensor rank should be: 1,2, or 3. Tensor rank is ", rank)
return
if rank==1:
dim0 = 1
dim1 = 1
dim2 = t.dim(0)
if rank==2:
dim0 = 1
dim1 = t.dim(0)
dim2 = t.dim(1)
if rank==3:
dim0 = t.dim(0)
dim1 = t.dim(1)
dim2 = t.dim(2)
var val:SIMD[dtype, 1]=0.0
for i in range(dim0):
if i==0 and rank==3:
print("[")
else:
if i>0:
print()
for j in range(dim1):
if rank!=1:
if j==0:
print_no_newline(" [")
else:
print_no_newline("\n ")
print_no_newline("[")
for k in range(dim2):
if rank==1:
val = t[k]
if rank==2:
val = t[j,k]
if rank==3:
val = t[i,j,k]
let int_str: String
if val > 0:
int_str = String(trunc(val).cast[DType.int32]())
else:
int_str = "-"+String(trunc(val).cast[DType.int32]())
val = -val
let float_str: String
float_str = String(mod(val,1))
let s = int_str+"."+float_str[2:prec+2]
if k==0:
print_no_newline(s)
else:
print_no_newline(" ",s)
print_no_newline("]")
if rank>1:
print_no_newline("]")
print()
if rank>2:
print("]")
print("Tensor shape:",t.shape().__str__(),", Tensor rank:",rank,",","DType:", dtype.__str__())
alias dtype = DType.float32
let tx = myTensor[dtype](5,10)
tx.print()
tx.mean().print()
output:
[[0.0850 0.8916 0.1896 0.3980 0.7435 0.5603 0.8095 0.5117 0.9950 0.9666]
[0.4260 0.6529 0.9615 0.8579 0.2940 0.4146 0.5148 0.7897 0.5442 0.0936]
[0.4322 0.8449 0.7728 0.1918 0.7803 0.1813 0.5791 0.3141 0.4119 0.9923]
[0.1639 0.3348 0.0762 0.1745 0.0372 0.4674 0.6741 0.0667 0.3897 0.1653]
[0.9908 0.8706 0.6726 0.5877 0.2550 0.5930 0.2717 0.2704 0.0959 0.6325]]
Tensor shape: 5x10 , Tensor rank: 2 , DType: float32
[[0.6151]
[0.5549]
[0.5501]
[0.2550]
[0.5240]]
Tensor shape: 5x1 , Tensor rank: 2 , DType: float32