diff --git a/builtins/edge/ds/checks.go b/builtins/edge/ds/checks.go
index fcc642ee..f5b6a54d 100644
--- a/builtins/edge/ds/checks.go
+++ b/builtins/edge/ds/checks.go
@@ -29,7 +29,7 @@ import (
 func RegisterChecks(logger *zerolog.Logger, fnName string, dr resolvers.DirectoryResolver) (*rego.Function, rego.Builtin1) {
 	return &rego.Function{
 			Name:    fnName,
-			Decl:    types.NewFunction(types.Args(types.A), types.B),
+			Decl:    types.NewFunction(types.Args(types.A), types.A),
 			Memoize: true,
 		},
 		func(bctx rego.BuiltinContext, op1 *ast.Term) (*ast.Term, error) {