Directional Color Shader

I wanted to replicate an effect that a friend made with directional light in Unity. The goal was to make a shader that colors an object based on its normal direction in world space. 

DirectionalColor.shader

Shader "Custom/DirectionalColor" {
    Properties {
        _ColorX ("ColorX", Color) = (1,0,0,1)
        _ColorY ("ColorY", Color) = (0,1,0,1)
        _ColorZ ("ColorZ", Color) = (0,0,1,1)
    }
    SubShader {
        Tags {
            "RenderType"="Opaque"
        }
        Pass {
            Name "FORWARD"
            Tags {
                "LightMode"="ForwardBase"
            }

            CGPROGRAM
            #pragma vertex vert
            #pragma fragment frag
            #define UNITY_PASS_FORWARDBASE
            #include "UnityCG.cginc"
            #pragma multi_compile_fwdbase_fullshadows
            #pragma exclude_renderers gles3 metal d3d11_9x xbox360 xboxone ps3 ps4 psp2
            #pragma target 3.0

            uniform float4 _ColorX;
            uniform float4 _ColorY;
            uniform float4 _ColorZ;

            struct VertexInput {
                float4 vertex : POSITION;
                float3 normal : NORMAL;
            };

            struct VertexOutput {
                float4 pos : SV_POSITION;
                float3 normalDir : TEXCOORD0;
            };

            VertexOutput vert (VertexInput v) {
                VertexOutput o = (VertexOutput)0;
                o.normalDir = UnityObjectToWorldNormal(v.normal);
                o.pos = mul(UNITY_MATRIX_MVP, v.vertex );
                return o;
            }

            float4 frag(VertexOutput i) : COLOR {
                i.normalDir = normalize(i.normalDir);
                float3 normalDirection = i.normalDir;
                float3 squared = (i.normalDir*i.normalDir);
                float3 finalColor = (squared.r*_ColorX.rgb + squared.g*_ColorY.rgb + squared.b*_ColorZ.rgb);
                return fixed4(finalColor,1);
            }
            ENDCG
        }
    }
    FallBack "Diffuse"
}