Updating Beta branch from main

README.md

@@ -177,8 +177,8 @@ train_args = TrainArguments(
     num_train_epochs=3,
     learning_rate=1e-5,
     learning_rate_end=1e-6,
-    optimizer='lion',  # 'adamw', 'lion', 'adafactor','warm_up_cosine' are supported
-    scheduler='linear',  # 'linear' or 'cosine' or 'none'
+    optimizer='lion',  # 'adamw', 'lion', 'adafactor' are supported
+    scheduler='linear',  # 'linear','cosine', 'none' ,'warm_up_cosine' and 'warm_up_linear'  are supported
     weight_decay=0.01,
     total_batch_size=16,
     max_steps=None,  # None to let trainer Decide

docs/Mojo/Array.md

@@ -130,46 +130,67 @@ Takes DType as dynamic Input like `Array[DType.float32]`
 
 `fn __init__(inout self: Self, array_shape: ArrayShape):`
 
- - Description: Init Array From ArrayShape(Alloc Zero).
+- Description: Init Array From ArrayShape(Alloc Zero).
 
 `fn __init__(inout self: Self, A: Self, B: Self) -> None:`
 
- - Description: Init Array From Two other Arrays A and B For Matmul(Alloc One).
+- Description: Init Array From Two other Arrays A and B For Matmul(Alloc One).
 
 `fn __init__(inout self: Self, vl: VariadicList[Int]):`
 
- - Description: Init Array from VariadicList[Int](Alloc Zero).
+- Description: Init Array from VariadicList[Int](Alloc Zero).
 
 `fn __init__(inout self: Self, init: Bool, *dim: Int) -> None:`
 
- - Description: Init Array from Int Args(Depends on Passed Bool).
+- Description: Init Array from Int Args(Depends on Passed Bool).
 
 `fn __init__(inout self: Self, *dim: Int):`
 
- - Description: Init Array from Int Args(Alloc Zero).
+- Description: Init Array from Int Args(Alloc Zero).
 
 `fn __init__(inout self: Self, value: DynamicVector[FloatLiteral], shape: ArrayShape) -> None:`
 
- - Description: Init Array from ArrayShape and load data from DynamicVector[FloatLiteral](Alloc One).
+- Description: Init Array from ArrayShape and load data from DynamicVector[FloatLiteral](Alloc One).
 
 `fn __init__(inout self: Self, value: VariadicList[FloatLiteral], shape: ArrayShape) -> None:`
 
- - Description: Init Array from ArrayShape and load data from VariadicList[FloatLiteral](Alloc One).
+- Description: Init Array from ArrayShape and load data from VariadicList[FloatLiteral](Alloc One).
 
 `fn __init__(inout self: Self, pointer: DTypePointer[T], *dim: Int) -> None:`
 
- - Description: Init Array from IntArgs and load data from DTypePointer[T](Alloc One).
+- Description: Init Array from IntArgs and load data from DTypePointer[T](Alloc One).
+
+`fn __init__(inout self: Self, pointer: DTypePointer[T], *dim: Int) -> None:`
+
+- Description: Init Array from given data from DTypePointer[T](Alloc Zero).
 
 ### Alloc
 
 `fn alloc(inout self: Self) -> None:`
- - Description: Allocate or Init The Array.
+
+- Description: Allocate or Init The Array.
+
+`fn alloc(inout self: Self, fill:SIMD[T, 1]) -> None:`
+
+- Allocate or Init The Array and fill that with given fill number.
 
 ### Random
 
 `fn random(inout self: Self) -> None:`
 
- - Description: Randomize The Data if the Array is Allocated.
+- Description: Randomize The Data if the Array is Allocated.
+
+### Reshape and View
+
+View:
+
+* INOUT  `fn view(inout self, *dims: Int):`
+* View Change Shape totaly and don't care if the new shape fits or doesn't.
+
+Reshape:
+
+* INOUT  `fn reshape(inout self, *dims: Int):`
+* Reshape Change Shape totaly and check if the new shape fits or doesn't.
 
 ### Dim
 
@@ -203,7 +224,7 @@ fn load[
 
 ### Store Functions
 
-``` 
+```
 fn store[
     nelts: Int, off: Int
 ](self, index: InlinedFixedVector[off, Int], val: SIMD[T, nelts]) -> None:
@@ -225,6 +246,10 @@ fn store[
 
 `fn __getitem__(self, index: Int) -> SIMD[T, 1]:`
 
+`fn __getitem__(self, d1: Int, d2: Int, val:SIMD[T, 1]) raises->None:`
+
+`fn __getitem__(self, d1: Int, d2: Int, d3: Int, val:SIMD[T, 1]) raises->None:`
+
 ### `__setitem__` Functions
 
 `fn __setitem__(self, index: Int, val: SIMD[T, 1]) -> None:`
@@ -233,9 +258,16 @@ fn store[
 
 `fn __setitem__[off: Int](self, index: StaticIntTuple[off], val: SIMD[T, 1]):`
 
+`fn __setitem__(self, d1: Int, d2: Int, val:SIMD[T, 1]) raises->None:`
+
+`fn __setitem__(self, d1: Int, d2: Int, d3: Int, val:SIMD[T, 1]) raises->None:`
+
 ### Math Functions
 
 ```mojo
+
+fn argmax(self: Self, axis: Int = -1) -> Int:
+
 # cos
 
 fn cos(inout self: Self) -> Self:
@@ -379,4 +411,4 @@ fn acosh(inout self: Self, rt: Runtime) -> Self:
 # Fill the whole array with the val 
 
 fn fill(inout self: Self, val: SIMD[T, 1]) -> None:
-```
+```

docs/Python/JAXServer.md

@@ -94,7 +94,7 @@ you can use this function outside the class like this
 
 ```python
 for string, num_used_tokens in server.process(
-        'im an string',
+        'im a string',
         greedy=False,
         max_new_tokens=256  # or None to use Maximum numbers passed in Config
 ):

examples/training/causal-lm/gpt-j.py

@@ -102,7 +102,7 @@
 flags.DEFINE_string(
     name="scheduler",
     default='cosine',
-    help='which scheduler to use (available schedulers are cosine linear none warm_up_cosine)'
+    help='which scheduler to use (available schedulers are cosine linear none warm_up_cosine , warm_up_linear)'
 )
 
 flags.DEFINE_string(

lib/mojo/EasyDel/__init__.🔥

@@ -6,13 +6,32 @@ EasyDel Library [MOJO/EasyDel]
 from .array import (
     Array,
     ArrayShape,
+    # Array Module
     convert_numpy_to_easydel_array,
     matmul_shape,
+    matmul_single_row,
     matmul_2d,
     matmul,
+    T_AXIS_0_2_1,  # LOW LEVEL AND PARALLELIZED , VECTORIZED
+    T_AXIS_2_1_0,  # LOW LEVEL AND PARALLELIZED , VECTORIZED
+    T_AXIS_1_0_2,  # LOW LEVEL AND PARALLELIZED , VECTORIZED
+    T_2D,  # LOW LEVEL AND PARALLELIZED , VECTORIZED
+    CAT_3D_AXIS_2,  # LOW LEVEL AND PARALLELIZED , VECTORIZED
+    rotate_half,  # LOW LEVEL AND PARALLELIZED , VECTORIZED
+    arange,
+    # Array Utils
+    softmax,
+    silu,
+    relu,
+    leaky_relu,
+    sigmoid,
+    sample_array,
+    # Funcs
+    PointerOperation,
+    # ED Struct
 )
 from .in_out import File, BufReader
-from .tokenizer import Tokenizer, load_tokenizer, byte_pr_tokenizer_encoder
+from .tokenizer import Tokenizer, load_tokenizer, byte_pr_tokenizer_encoder, loop_sort
 from .utilities import (
     read_numerical_value,
     read_string_value,

lib/mojo/EasyDel/array/__init__.🔥

@@ -8,6 +8,16 @@ from .array_module import Array, ArrayShape
 from .array_utils import (
     convert_numpy_to_easydel_array,
     matmul_shape,
+    matmul_single_row,
     matmul_2d,
     matmul,
+    T_AXIS_0_2_1,  # LOW LEVEL AND PARALLELIZED , VECTORIZED
+    T_AXIS_2_1_0,  # LOW LEVEL AND PARALLELIZED , VECTORIZED
+    T_AXIS_1_0_2,  # LOW LEVEL AND PARALLELIZED , VECTORIZED
+    T_2D,  # LOW LEVEL AND PARALLELIZED , VECTORIZED
+    CAT_3D_AXIS_2,  # LOW LEVEL AND PARALLELIZED , VECTORIZED
+    rotate_half,  # LOW LEVEL AND PARALLELIZED , VECTORIZED
+    arange,
 )
+from .funcs import softmax, silu, relu, leaky_relu, sigmoid, sample_array
+from .ed_struct import PointerOperation