examples: mnist-rs: bump version of burn

Author: ivila

examples/mnist-rs/Makefile

@@ -20,10 +20,9 @@
 CROSS_COMPILE_HOST ?= aarch64-linux-gnu-
 CROSS_COMPILE_TA ?= aarch64-linux-gnu-
 TARGET_HOST ?= aarch64-unknown-linux-gnu
-TARGET_TA ?= aarch64-unknown-linux-gnu
+TARGET_TA := aarch64-unknown-none
 FEATURES ?=
 CARGO_FLAGS ?=
-
 .PHONY: host ta all clean
 
 all: toolchain host ta

examples/mnist-rs/ta/Cargo.toml

@@ -39,7 +39,7 @@ optee-utee-build = { path = "../../../crates/optee-utee-build" }
 proto = { path = "../proto" }
 
 bytemuck = { version = "1.21.0", features = ["min_const_generics"] }
-burn = { version = "0.17", default-features = false, features = ["ndarray", "autodiff"] }
+burn = { version = "0.21", default-features = false, features = ["ndarray", "autodiff"] }
 spin = "0.9.8"
 serde = { version = "1.0.218", default-features = false, features = ["derive"] }
 serde_json = { version = "1.0.139", default-features = false, features = ["alloc"] }

examples/mnist-rs/ta/inference/Makefile

@@ -18,13 +18,14 @@
 UUID ?= $(shell cat "./uuid.txt")
 NAME := inference
 
-TARGET ?= aarch64-unknown-linux-gnu
+TARGET ?= aarch64-unknown-none
 CROSS_COMPILE ?= aarch64-linux-gnu-
 OBJCOPY := $(CROSS_COMPILE)objcopy
 # Configure the linker to use GCC, which works on both cross-compilation and ARM machines
 LINKER_CFG := target.$(TARGET).linker=\"$(CROSS_COMPILE)gcc\"
 # fix for the error: "unwinding panics are not supported without std" reported by clippy
-RUSTFLAGS := -C panic=abort
+RUSTFLAGS := -C panic=abort -C relocation-model=pic
+CARGO_FLAGS ?=
 
 TA_SIGN_KEY ?= $(TA_DEV_KIT_DIR)/keys/default_ta.pem
 SIGN := $(TA_DEV_KIT_DIR)/scripts/sign_encrypt.py
@@ -34,10 +35,10 @@ all: clippy ta strip sign
 
 clippy:
 	@cargo fmt
-	@RUSTFLAGS="$(RUSTFLAGS)" cargo clippy --target $(TARGET) -- -D warnings -D clippy::unwrap_used -D clippy::expect_used -D clippy::panic
+	@RUSTFLAGS="$(RUSTFLAGS)" cargo clippy $(CARGO_FLAGS) --target $(TARGET) -- -D warnings -D clippy::unwrap_used -D clippy::expect_used -D clippy::panic
 
 ta: clippy
-	@RUSTFLAGS="$(RUSTFLAGS)" cargo build --target $(TARGET) --release --config $(LINKER_CFG)
+	@RUSTFLAGS="$(RUSTFLAGS)" cargo build $(CARGO_FLAGS) --target $(TARGET) --release --config $(LINKER_CFG)
 
 strip: ta
 	@$(OBJCOPY) --strip-unneeded $(OUT_DIR)/$(NAME) $(OUT_DIR)/stripped_$(NAME)

examples/mnist-rs/ta/train/Makefile

@@ -18,13 +18,14 @@
 UUID ?= $(shell cat "./uuid.txt")
 NAME := train
 
-TARGET ?= aarch64-unknown-linux-gnu
+TARGET ?= aarch64-unknown-none
 CROSS_COMPILE ?= aarch64-linux-gnu-
 OBJCOPY := $(CROSS_COMPILE)objcopy
 # Configure the linker to use GCC, which works on both cross-compilation and ARM machines
 LINKER_CFG := target.$(TARGET).linker=\"$(CROSS_COMPILE)gcc\"
 # fix for the error: "unwinding panics are not supported without std" reported by clippy
-RUSTFLAGS := -C panic=abort
+RUSTFLAGS := -C panic=abort -C relocation-model=pic
+CARGO_FLAGS ?=
 
 TA_SIGN_KEY ?= $(TA_DEV_KIT_DIR)/keys/default_ta.pem
 SIGN := $(TA_DEV_KIT_DIR)/scripts/sign_encrypt.py
@@ -34,10 +35,10 @@ all: clippy ta strip sign
 
 clippy:
 	@cargo fmt
-	@RUSTFLAGS="$(RUSTFLAGS)" cargo clippy --target $(TARGET) -- -D warnings -D clippy::unwrap_used -D clippy::expect_used -D clippy::panic
+	@RUSTFLAGS="$(RUSTFLAGS)" cargo clippy $(CARGO_FLAGS) --target $(TARGET) -- -D warnings -D clippy::unwrap_used -D clippy::expect_used -D clippy::panic
 
 ta: clippy
-	@RUSTFLAGS="$(RUSTFLAGS)" cargo build --target $(TARGET) --release --config $(LINKER_CFG)
+	@RUSTFLAGS="$(RUSTFLAGS)" cargo build $(CARGO_FLAGS) --target $(TARGET) --release --config $(LINKER_CFG)
 
 strip: ta
 	@$(OBJCOPY) --strip-unneeded $(OUT_DIR)/$(NAME) $(OUT_DIR)/stripped_$(NAME)

examples/mnist-rs/ta/train/src/trainer.rs

@@ -38,7 +38,7 @@ impl<B: AutodiffBackend> Trainer<B> {
     pub fn new(device: B::Device, lr: f64) -> Self {
         let mut seed = [0_u8; 8];
         optee_utee::Random::generate(seed.as_mut_slice());
-        B::seed(u64::from_le_bytes(seed));
+        B::seed(&device, u64::from_le_bytes(seed));
 
         Self {
             optim: AdamConfig::new().init(),
@@ -106,7 +106,7 @@ impl<B: AutodiffBackend> Trainer<B> {
 // to https://github.com/tracel-ai/burn/blob/v0.16.0/crates/burn-no-std-test for
 // details.
 fn accuracy<B: Backend>(output: Tensor<B, 2>, targets: Tensor<B, 1, Int>) -> f32 {
-    let predictions = output.argmax(1).squeeze(1);
+    let predictions = output.argmax(1).squeeze();
     let num_predictions: usize = targets.dims().iter().product();
     let num_corrects = predictions.equal(targets).int().sum().into_scalar();