# ZygoteRules dispatch on traits?

**URL:** <https://discourse.julialang.org/t/zygoterules-dispatch-on-traits/64869>\
**Category:** Machine Learning\
**Created:** [July 18, 2021, 6:38pm UTC](https://discourse.julialang.org/t/zygoterules-dispatch-on-traits/64869 "2021-07-18T18:38:46Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![lamorton](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lamorton/32/21115_2.png) [@lamorton](https://discourse.julialang.org/u/lamorton)\
**Post date:** [July 18, 2021, 6:38pm UTC](https://discourse.julialang.org/t/zygoterules-dispatch-on-traits/64869/1 "2021-07-18T18:38:46Z")

</div>

Is it possible to specialize a ZygoteRule using a Holy trait? I’m not sure how I would go about doing this.

---

<div class="post-metadata">

**Author:** ![lamorton](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lamorton/32/21115_2.png) [@lamorton](https://discourse.julialang.org/u/lamorton)\
**Post date:** [July 19, 2021, 5:42pm UTC](https://discourse.julialang.org/t/zygoterules-dispatch-on-traits/64869/2 "2021-07-19T17:42:37Z")

</div>

OK, here’s what I have come up with so far (doesn’t work, but shows what I’m trying to do):

```julia
using Zygote
abstract type SymbolicTrait end
struct IsSymbolic <: SymbolicTrait end
struct NotSymbolic <: SymbolicTrait end
function SymbolicTrait(x)
    if typeof(x) <: Int
        IsSymbolic()
    else
        NotSymbolic()
    end
end

function bloop(x)
    x^2
end

Zygote.@adjoint function bloop(x)
     bloop_adjoint(SymbolicTrait(x),x) #How do I actually refer to the adjoint defined by the @adjoint macro?
end

Zygote.@adjoint function bloop(::IsSymbolic, x)
    function bloop_adjoint(d)
        2*d
        print("Symbolic")
    end
    (x^2,bloop_adjoint)
end
Zygote.@adjoint function bloop(::NotSymbolic, x)
    function bloop_adjoint(d)
        2*d
        print("NonSymbolic")
    end
    (x^2,bloop_adjoint)
end

```

I’m using Int/Float to stand in for `Symbolic`/numeric types for simplicity of the example.

```julia
julia> Zygote.gradient(bloop,1)
ERROR: UndefVarError: bloop_adjoint not defined
Stacktrace:
 [1] adjoint
   @ ./REPL[7]:2 [inlined]
 [2] _pullback( __context__ ::Zygote.Context, 276::typeof(bloop), x::Int64)
   @ Main ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:57
 [3] _pullback(f::Function, args::Int64)
   @ Zygote ~/.julia/packages/Zygote/0da6K/src/compiler/interface.jl:34
 [4] pullback(f::Function, args::Int64)
   @ Zygote ~/.julia/packages/Zygote/0da6K/src/compiler/interface.jl:40
 [5] gradient(f::Function, args::Int64)
   @ Zygote ~/.julia/packages/Zygote/0da6K/src/compiler/interface.jl:58
 [6] top-level scope
   @ REPL[10]:100:

```

---

<div class="post-metadata">

**Author:** ![sbuercklin](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sbuercklin/32/15728_2.png) [@sbuercklin](https://discourse.julialang.org/u/sbuercklin)\
**Post date:** [July 19, 2021, 5:52pm UTC](https://discourse.julialang.org/t/zygoterules-dispatch-on-traits/64869/3 "2021-07-19T17:52:05Z")

</div>

The function you call within `bloop(x)` isn’t visible within that `bloop` method. Instead, try returning `bloop(SymbolicTrait(x),x)`.

You could also at that point remove the `@adjoint` annotations from your other two `bloop` methods, as they’re only interfaced with internally from your first `bloop` method; Zygote will only deal with the returned adjoint and primal from `bloop(x)`, which wraps the other two methods.

Lastly, unless you’re specifically overloading a method which has an existing adjoint, you might consider moving to `ChainRulesCore` to define your `rrule`s, rather than sticking with `ZygoteRules.@adjoint`. I believe the plan is to migrate away from `ZygoteRules` and over to `ChainRules` exclusively in the future.