WIP

crates/burn-tensor/src/lib.rs

@@ -11,6 +11,9 @@ extern crate derive_new;
 
 extern crate alloc;
 
+/// Server module.
+pub mod server;
+
 mod tensor;
 
 /// Burn Tensor representaton

crates/burn-tensor/src/server/backend.rs

@@ -0,0 +1,102 @@
+use core::marker::PhantomData;
+
+use crate::{
+    backend::{Backend, BackendBridge, DeviceOps},
+    quantization::QTensorPrimitive,
+    repr::{OperationDescription, TensorDescription},
+};
+
+pub struct Server<B: ServerBackend> {
+    r: PhantomData<B>,
+}
+
+pub trait ServerBackend: Send + Sync + 'static {
+    type Runtime: ServerRuntime;
+}
+
+pub trait ServerRuntime {
+    type Client: ServerClient;
+    type Device: DeviceOps;
+    type Bridge;
+}
+
+pub struct ServerTensor<R: ServerRuntime> {
+    desc: TensorDescription,
+    client: R::Client,
+}
+
+impl<B: ServerBackend> core::fmt::Debug for Server<B> {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.write_fmt(format_args!("server"))
+    }
+}
+
+impl<R: ServerRuntime> core::fmt::Debug for ServerTensor<R> {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.write_fmt(format_args!("tensor"))
+    }
+}
+
+impl<R: ServerRuntime> Clone for ServerTensor<R> {
+    fn clone(&self) -> Self {
+        Self {
+            desc: self.desc.clone(),
+            client: self.client.clone(),
+        }
+    }
+}
+
+impl<B: ServerBackend> Clone for Server<B> {
+    fn clone(&self) -> Self {
+        Self { r: PhantomData }
+    }
+}
+
+impl<B: ServerBackend> Default for Server<B> {
+    fn default() -> Self {
+        Self { r: PhantomData }
+    }
+}
+
+impl<R: ServerRuntime> QTensorPrimitive for ServerTensor<R> {
+    fn scheme(&self) -> &crate::quantization::QuantizationScheme {
+        todo!()
+    }
+
+    fn strategy(&self) -> crate::quantization::QuantizationStrategy {
+        todo!()
+    }
+}
+
+pub trait ServerClient: Clone + Send + Sync {
+    fn execute(&self, op: OperationDescription);
+}
+
+impl<B: ServerBackend> Backend for Server<B>
+where
+    <B::Runtime as ServerRuntime>::Bridge: BackendBridge<Self> + 'static,
+{
+    type Device = <B::Runtime as ServerRuntime>::Device;
+
+    type FullPrecisionBridge = <B::Runtime as ServerRuntime>::Bridge;
+
+    type FloatTensorPrimitive<const D: usize> = ServerTensor<B::Runtime>;
+
+    type FloatElem = f32;
+
+    type IntTensorPrimitive<const D: usize> = ServerTensor<B::Runtime>;
+
+    type IntElem = i32;
+
+    type BoolTensorPrimitive<const D: usize> = ServerTensor<B::Runtime>;
+
+    type QuantizedTensorPrimitive<const D: usize> = ServerTensor<B::Runtime>;
+
+    fn name() -> String {
+        todo!()
+    }
+
+    fn seed(seed: u64) {
+        todo!()
+    }
+}

crates/burn-tensor/src/server/mod.rs

@@ -0,0 +1,4 @@
+mod backend;
+mod ops;
+
+pub use backend::*;

crates/burn-tensor/src/server/ops/mod.rs

@@ -0,0 +1,6 @@
+mod op_activation;
+mod op_bool;
+mod op_float;
+mod op_int;
+mod op_module;
+mod op_qfloat;

crates/burn-tensor/src/server/ops/op_activation.rs

@@ -0,0 +1,4 @@
+use crate::server::Server;
+use crate::{ops::ActivationOps, server::ServerBackend};
+
+impl<B: ServerBackend> ActivationOps<Self> for Server<B> {}

crates/burn-tensor/src/server/ops/op_bool.rs

@@ -0,0 +1,119 @@
+use crate::{
+    ops::BoolTensorOps,
+    server::{Server, ServerBackend},
+};
+
+impl<B: ServerBackend> BoolTensorOps<Self> for Server<B> {
+    fn bool_empty<const D: usize>(
+        shape: crate::Shape<D>,
+        device: &crate::Device<Self>,
+    ) -> crate::ops::BoolTensor<Self, D> {
+        todo!()
+    }
+
+    fn bool_shape<const D: usize>(tensor: &crate::ops::BoolTensor<Self, D>) -> crate::Shape<D> {
+        todo!()
+    }
+
+    fn bool_into_data<const D: usize>(
+        tensor: crate::ops::BoolTensor<Self, D>,
+    ) -> impl core::future::Future<Output = crate::TensorData> + Send {
+        todo!()
+    }
+
+    fn bool_from_data<const D: usize>(
+        data: crate::TensorData,
+        device: &crate::Device<Self>,
+    ) -> crate::ops::BoolTensor<Self, D> {
+        todo!()
+    }
+
+    fn bool_into_int<const D: usize>(
+        tensor: crate::ops::BoolTensor<Self, D>,
+    ) -> crate::ops::IntTensor<Self, D> {
+        todo!()
+    }
+
+    fn bool_into_float<const D: usize>(
+        tensor: crate::ops::BoolTensor<Self, D>,
+    ) -> crate::ops::FloatTensor<Self, D> {
+        todo!()
+    }
+
+    fn bool_device<const D: usize>(
+        tensor: &crate::ops::BoolTensor<Self, D>,
+    ) -> crate::Device<Self> {
+        todo!()
+    }
+
+    fn bool_to_device<const D: usize>(
+        tensor: crate::ops::BoolTensor<Self, D>,
+        device: &crate::Device<Self>,
+    ) -> crate::ops::BoolTensor<Self, D> {
+        todo!()
+    }
+
+    fn bool_reshape<const D1: usize, const D2: usize>(
+        tensor: crate::ops::BoolTensor<Self, D1>,
+        shape: crate::Shape<D2>,
+    ) -> crate::ops::BoolTensor<Self, D2> {
+        todo!()
+    }
+
+    fn bool_slice<const D1: usize, const D2: usize>(
+        tensor: crate::ops::BoolTensor<Self, D1>,
+        ranges: [core::ops::Range<usize>; D2],
+    ) -> crate::ops::BoolTensor<Self, D1> {
+        todo!()
+    }
+
+    fn bool_slice_assign<const D1: usize, const D2: usize>(
+        tensor: crate::ops::BoolTensor<Self, D1>,
+        ranges: [core::ops::Range<usize>; D2],
+        value: crate::ops::BoolTensor<Self, D1>,
+    ) -> crate::ops::BoolTensor<Self, D1> {
+        todo!()
+    }
+
+    fn bool_equal<const D: usize>(
+        lhs: crate::ops::BoolTensor<Self, D>,
+        rhs: crate::ops::BoolTensor<Self, D>,
+    ) -> crate::ops::BoolTensor<Self, D> {
+        todo!()
+    }
+
+    fn bool_not<const D: usize>(
+        tensor: crate::ops::BoolTensor<Self, D>,
+    ) -> crate::ops::BoolTensor<Self, D> {
+        todo!()
+    }
+
+    fn bool_swap_dims<const D: usize>(
+        tensor: crate::ops::BoolTensor<Self, D>,
+        dim1: usize,
+        dim2: usize,
+    ) -> crate::ops::BoolTensor<Self, D> {
+        todo!()
+    }
+
+    fn bool_permute<const D: usize>(
+        tensor: crate::ops::BoolTensor<Self, D>,
+        axes: [usize; D],
+    ) -> crate::ops::BoolTensor<Self, D> {
+        todo!()
+    }
+
+    fn bool_flip<const D: usize>(
+        tensor: crate::ops::BoolTensor<Self, D>,
+        axes: &[usize],
+    ) -> crate::ops::BoolTensor<Self, D> {
+        todo!()
+    }
+
+    fn bool_expand<const D1: usize, const D2: usize>(
+        tensor: crate::ops::BoolTensor<Self, D1>,
+        shape: crate::Shape<D2>,
+    ) -> crate::ops::BoolTensor<Self, D2> {
+        todo!()
+    }
+}

crates/burn-tensor/src/server/ops/op_float.rs

@@ -0,0 +1,4 @@
+use crate::server::Server;
+use crate::{ops::FloatTensorOps, server::ServerBackend};
+
+impl<B: ServerBackend> FloatTensorOps<Self> for Server<B> {}

crates/burn-tensor/src/server/ops/op_int.rs

@@ -0,0 +1,4 @@
+use crate::server::Server;
+use crate::{ops::IntTensorOps, server::ServerBackend};
+
+impl<B: ServerBackend> IntTensorOps<Self> for Server<B> {}

crates/burn-tensor/src/server/ops/op_module.rs

@@ -0,0 +1,7 @@
+use crate::server::Server;
+use crate::{
+    ops::{ConvOptions, ConvTransposeOptions, FloatTensor, ModuleOps},
+    server::ServerBackend,
+};
+
+impl<B: ServerBackend> ModuleOps<Self> for Server<B> {}

crates/burn-tensor/src/server/ops/op_qfloat.rs

@@ -0,0 +1,10 @@
+use core::future::Future;
+
+use crate::{
+    ops::{FloatTensor, QTensorOps, QuantizedTensor},
+    quantization::{QuantizationParametersPrimitive, QuantizationScheme},
+    server::{Server, ServerBackend},
+    Device, TensorData,
+};
+
+impl<B: ServerBackend> QTensorOps<Self> for Server<B> {}

crates/burn-tensor/src/tensor/backend/base.rs

@@ -60,8 +60,8 @@ pub trait Backend:
     + ActivationOps<Self>
     + QTensorOps<Self>
     + Clone
-    + Sized
     + Default
+    + Sized
     + Send
     + Sync
     + core::fmt::Debug