From 4981557ba0a4e5d20e9c4dc770af85d488908925 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20S=C3=A1nchez=20Ram=C3=ADrez?=
 <15837247+mofeing@users.noreply.github.com>
Date: Wed, 27 Nov 2024 22:27:48 +0100
Subject: [PATCH] fix `YaoBlocks.mat` for rotation gates with `TracedRNumber`
 parameters (#306)

* fix `YaoBlocks.mat` for rotation gates with `TracedRNumber` parameters

* Update ext/ReactantYaoBlocksExt.jl
---
 Project.toml                |  2 ++
 ext/ReactantYaoBlocksExt.jl | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+)
 create mode 100644 ext/ReactantYaoBlocksExt.jl

diff --git a/Project.toml b/Project.toml
index 631ac23c4..170f0efe5 100644
--- a/Project.toml
+++ b/Project.toml
@@ -23,12 +23,14 @@ AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+YaoBlocks = "418bc28f-b43b-5e0b-a6e7-61bbc1a2c1df"
 
 [extensions]
 ReactantAbstractFFTsExt = "AbstractFFTs"
 ReactantArrayInterfaceExt = "ArrayInterface"
 ReactantNNlibExt = "NNlib"
 ReactantStatisticsExt = "Statistics"
+ReactantYaoBlocksExt = "YaoBlocks"
 
 [compat]
 AbstractFFTs = "1.5"
diff --git a/ext/ReactantYaoBlocksExt.jl b/ext/ReactantYaoBlocksExt.jl
new file mode 100644
index 000000000..a19a105ad
--- /dev/null
+++ b/ext/ReactantYaoBlocksExt.jl
@@ -0,0 +1,37 @@
+module ReactantYaoBlocksExt
+
+using Reactant
+using YaoBlocks
+
+function YaoBlocks.mat(
+    ::Type{T}, R::RotationGate{D,Reactant.TracedRNumber{S},<:XGate}
+) where {D,T,S}
+    M = Reactant.broadcast_to_size(zero(T), (2, 2))
+    M[1, 1] = cos(R.theta / 2)
+    M[2, 2] = cos(R.theta / 2)
+    M[1, 2] = -im * sin(R.theta / 2)
+    M[2, 1] = -im * sin(R.theta / 2)
+    return M
+end
+
+function YaoBlocks.mat(
+    ::Type{T}, R::RotationGate{D,Reactant.TracedRNumber{S},<:YGate}
+) where {D,T,S}
+    M = Reactant.broadcast_to_size(zero(T), (2, 2))
+    M[1, 1] = cos(R.theta / 2)
+    M[2, 2] = cos(R.theta / 2)
+    M[1, 2] = -sin(R.theta / 2)
+    M[2, 1] = sin(R.theta / 2)
+    return M
+end
+
+function YaoBlocks.mat(
+    ::Type{T}, R::RotationGate{D,Reactant.TracedRNumber{S},<:ZGate}
+) where {D,T,S}
+    M = Reactant.broadcast_to_size(zero(T), (2, 2))
+    M[1, 1] = exp(-im * R.theta / 2)
+    M[2, 2] = exp(im * R.theta / 2)
+    return M
+end
+
+end